diff --git a/lib/shared/src/sourcegraph-api/completions/CompletionsResponseBuilder.test.ts b/lib/shared/src/sourcegraph-api/completions/CompletionsResponseBuilder.test.ts new file mode 100644 index 000000000000..4b4c745c37e7 --- /dev/null +++ b/lib/shared/src/sourcegraph-api/completions/CompletionsResponseBuilder.test.ts @@ -0,0 +1,98 @@ +import { describe, expect, it } from 'vitest' + +import { CompletionsResponseBuilder } from './CompletionsResponseBuilder' + +interface CompletionsResponseTestCase { + name: string + url: string + steps: { + deltaText?: string + completion?: string + thinking?: string + expected: string + }[] +} + +describe('CompletionsResponseBuilder', () => { + const testCases: CompletionsResponseTestCase[] = [ + { + name: 'API v1 - Full replacement mode', + url: 'https://sourcegraph.com/.api/completions/stream?api-version=1', + steps: [ + { + completion: 'Direct', + expected: 'Direct', + }, + { + completion: ' Response', + expected: ' Response', + }, + ], + }, + { + name: 'API v2 - Incremental mode', + url: 'https://sourcegraph.com/.api/completions/stream?api-version=2', + steps: [ + { + deltaText: undefined, + expected: '', + }, + { + deltaText: undefined, + expected: '', + }, + { + deltaText: 'Starting response', + expected: 'Starting response', + }, + ], + }, + { + name: 'API v2 - Incremental mode with thinking steps', + url: 'https://sourcegraph.com/.api/completions/stream?api-version=2', + steps: [ + { + thinking: 'Analyzing...', + expected: 'Analyzing...\n', + }, + { + thinking: 'Refining...', + expected: 'Analyzing...Refining...\n', + }, + { + deltaText: 'Better response', + expected: 'Analyzing...Refining...\nBetter response', + }, + ], + }, + { + name: 'API v8 - Incremental mode with thinking steps', + url: 'https://sourcegraph.com/.api/completions/stream?api-version=8', + steps: [ + { + thinking: 'Step 1...', + deltaText: 'Hello', + expected: 'Step 1...\nHello', + }, + { + thinking: 'Step 2...', + deltaText: ' World', + expected: 'Step 1...Step 2...\nHello World', + }, + ], + }, + ] + + for (const testCase of testCases) { + describe(testCase.name, () => { + it('processes completion steps correctly', () => { + const builder = CompletionsResponseBuilder.fromUrl(testCase.url) + for (const step of testCase.steps) { + builder.nextThinking(step.thinking ?? undefined) + const result = builder.nextCompletion(step.completion, step.deltaText) + expect(result).toBe(step.expected) + } + }) + }) + } +}) diff --git a/lib/shared/src/sourcegraph-api/completions/CompletionsResponseBuilder.ts b/lib/shared/src/sourcegraph-api/completions/CompletionsResponseBuilder.ts index cd01863400f2..46e8c5c1dd03 100644 --- a/lib/shared/src/sourcegraph-api/completions/CompletionsResponseBuilder.ts +++ b/lib/shared/src/sourcegraph-api/completions/CompletionsResponseBuilder.ts @@ -8,12 +8,39 @@ export class CompletionsResponseBuilder { public totalCompletion = '' constructor(public readonly apiVersion: number) {} public nextCompletion(completion: string | undefined, deltaText: string | undefined): string { + const thinkingText = this.getThinkingText() if (this.apiVersion >= 2) { this.totalCompletion += deltaText ?? '' - return this.totalCompletion + } else { + this.totalCompletion = completion ?? '' } - this.totalCompletion = completion ?? '' - return this.totalCompletion + return thinkingText + this.totalCompletion + } + + private readonly thinkingBuffer: string[] = [] + /** + * Processes and accumulates thinking steps during the completion stream. + * Thinking steps must start at the beginning of completion and are enclosed in tags. + * When the completion starts streaming, the previous tag is closed. + * + * @param deltaThinking - The incremental thinking text to be added + * @returns The formatted thinking text wrapped in XML tags + */ + public nextThinking(deltaThinking?: string): string { + if (deltaThinking) { + this.thinkingBuffer.push(deltaThinking) + } + return this.getThinkingText() + } + /** + * Generates the formatted thinking text by combining all thinking steps. + * Wraps the combined thinking text in tags and adds a newline if content exists. + * + * @returns Formatted thinking text with XML tags, or empty string if no thinking steps exist + */ + private getThinkingText(): string { + const thinking = this.thinkingBuffer.join('') + return thinking ? `${thinking}\n` : '' } public static fromUrl(url: string): CompletionsResponseBuilder { diff --git a/lib/shared/src/sourcegraph-api/completions/parse.ts b/lib/shared/src/sourcegraph-api/completions/parse.ts index 4f0e04e82309..8eca74000e23 100644 --- a/lib/shared/src/sourcegraph-api/completions/parse.ts +++ b/lib/shared/src/sourcegraph-api/completions/parse.ts @@ -33,6 +33,7 @@ function parseJSON(data: string): T | Error { export interface CompletionData { completion?: string deltaText?: string + delta_thinking?: string stopReason?: string } @@ -56,6 +57,9 @@ function parseEventData( if (isError(data)) { return data } + // Process the delta_thinking and deltaText separately. + // The thinking text will be added to the completion text. + builder.nextThinking(data.delta_thinking) // Internally, don't handle delta text yet and there's limited value // in passing around deltas anyways so we concatenate them here. const completion = builder.nextCompletion(data.completion, data.deltaText) diff --git a/lib/shared/src/sourcegraph-api/completions/types.ts b/lib/shared/src/sourcegraph-api/completions/types.ts index ae5e9654e2f4..5a61bb57cc64 100644 --- a/lib/shared/src/sourcegraph-api/completions/types.ts +++ b/lib/shared/src/sourcegraph-api/completions/types.ts @@ -40,6 +40,7 @@ export interface PromptTokensDetails { export interface CompletionResponse { completion: string + thinking?: string stopReason?: string } diff --git a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx index ff5600a4ebdc..4d3c0ef0b299 100644 --- a/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx +++ b/vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx @@ -3,6 +3,7 @@ import type React from 'react' import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import { clsx } from 'clsx' +import { LoaderIcon, PlusIcon } from 'lucide-react' import type { FixupTaskID } from '../../../src/non-stop/FixupTask' import { CodyTaskState } from '../../../src/non-stop/state' import { type ClientActionListener, useClientActionListener } from '../../client/clientState' @@ -12,7 +13,7 @@ import type { PriorHumanMessageInfo } from '../cells/messageCell/assistant/Assis import styles from './ChatMessageContent.module.css' import { GuardrailsStatusController } from './GuardRailStatusController' import { createButtons, createButtonsExperimentalUI } from './create-buttons' -import { getCodeBlockId, getFileName } from './utils' +import { extractThinkContent, getCodeBlockId, getFileName } from './utils' export interface CodeBlockActionsProps { copyButtonOnSubmit: (text: string, event?: 'Keydown' | 'Button') => void @@ -204,10 +205,43 @@ export const ChatMessageContent: React.FunctionComponent extractThinkContent(displayMarkdown), + [displayMarkdown] + ) + return (
+ {thinkContent.length > 0 && ( +
+ + {isThinking ? ( + + ) : ( + + )} + + {isThinking ? 'Thinking...' : 'Thought Process'} + + +
+ {thinkContent} +
+
+ )} - {displayMarkdown} + {displayContent}
) diff --git a/vscode/webviews/chat/ChatMessageContent/utils.test.ts b/vscode/webviews/chat/ChatMessageContent/utils.test.ts new file mode 100644 index 000000000000..92350ce2cac5 --- /dev/null +++ b/vscode/webviews/chat/ChatMessageContent/utils.test.ts @@ -0,0 +1,70 @@ +import { describe, expect, it } from 'vitest' +import { extractThinkContent } from './utils' + +describe('extractThinkContent', () => { + it('extracts content from complete think tags at the start', () => { + const input = 'Planning stepsHere is the code' + const result = extractThinkContent(input) + + expect(result).toEqual({ + displayContent: 'Here is the code', + thinkContent: 'Planning steps', + isThinking: false, + }) + }) + + it('ignores think tags that do not start at the beginning', () => { + const input = 'Code hereStep 2More code' + const result = extractThinkContent(input) + + expect(result).toEqual({ + displayContent: 'Code hereStep 2More code', + thinkContent: '', + isThinking: false, + }) + }) + + it('handles unclosed think tag at the start', () => { + const input = 'Incomplete thought' + const result = extractThinkContent(input) + + expect(result).toEqual({ + displayContent: '', + thinkContent: 'Incomplete thought', + isThinking: true, + }) + }) + + it('ignores unclosed think tag not at the start', () => { + const input = 'MiddleIncomplete' + const result = extractThinkContent(input) + + expect(result).toEqual({ + displayContent: 'MiddleIncomplete', + thinkContent: '', + isThinking: false, + }) + }) + + it('returns empty think content for input without think tags', () => { + const input = 'Regular content without think tags' + const result = extractThinkContent(input) + + expect(result).toEqual({ + displayContent: 'Regular content without think tags', + thinkContent: '', + isThinking: false, + }) + }) + + it('keeps isThinking true when think tag is closed but no content follows', () => { + const input = 'Completed thought' + const result = extractThinkContent(input) + + expect(result).toEqual({ + displayContent: '', + thinkContent: 'Completed thought', + isThinking: true, + }) + }) +}) diff --git a/vscode/webviews/chat/ChatMessageContent/utils.ts b/vscode/webviews/chat/ChatMessageContent/utils.ts index ba45399f7bb2..bf3e409e224c 100644 --- a/vscode/webviews/chat/ChatMessageContent/utils.ts +++ b/vscode/webviews/chat/ChatMessageContent/utils.ts @@ -11,3 +11,59 @@ export function getCodeBlockId(contents: string, fileName?: string): string { } return SHA256(input).toString() } + +interface StreamingContent { + displayContent: string + thinkContent: string + isThinking: boolean +} + +/** + * Extracts content enclosed in `` tags from the beginning of a string. + * This function processes text that may contain special thinking content markers. + * + * @param content - The input string that may contain thinking content + * @returns A StreamingContent object with three properties: + * - displayContent: The portion of the input string that should be displayed to the user + * (excludes content in think tags at the start) + * - thinkContent: The content found inside the think tags, if any + * - isThinking: A boolean indicating whether we're in "thinking" mode: + * true when there's either an unclosed think tag or + * a complete think tag with no content after it + * + * Note: Only think tags at the start of the content are processed. + * Think tags appearing later in the content are left as-is in displayContent. + */ +const lengthOfThinkTag = ''.length +export function extractThinkContent(content: string): StreamingContent { + // Match think tags at the start of the content + const thinkRegex = /^([\s\S]*?)<\/think>/ + const match = content.match(thinkRegex) + + // Check if content starts with a think tag + const startsWithThink = content.startsWith('') + + let thinkContent = '' + let displayContent = content + let isThinking = false + + if (match) { + // We found a complete think tag at the start + thinkContent = match[1].trim() + displayContent = content.slice(match[0].length) + + // If there's no content after the think tag, we're still in thinking mode + isThinking = displayContent.trim() === '' + } else if (startsWithThink) { + // We have an unclosed think tag at the start + thinkContent = content.slice(lengthOfThinkTag) // length of '' + displayContent = '' + isThinking = true + } + + return { + displayContent, + thinkContent, + isThinking, + } +} diff --git a/vscode/webviews/components/MarkdownFromCody.tsx b/vscode/webviews/components/MarkdownFromCody.tsx index fb13b1c3c284..4a8829897717 100644 --- a/vscode/webviews/components/MarkdownFromCody.tsx +++ b/vscode/webviews/components/MarkdownFromCody.tsx @@ -53,6 +53,7 @@ const ALLOWED_ELEMENTS = [ 'h5', 'h6', 'br', + 'think', ] function defaultUrlProcessor(url: string): string { diff --git a/vscode/webviews/tailwind.config.mjs b/vscode/webviews/tailwind.config.mjs index c53da6e65caf..05519b93146a 100644 --- a/vscode/webviews/tailwind.config.mjs +++ b/vscode/webviews/tailwind.config.mjs @@ -104,6 +104,10 @@ export default { background: 'var(--vscode-statusBarItem-offlineBackground)', foreground: 'var(--vscode-statusBarItem-offlineForeground)', }, + code: { + background: 'var(--code-background)', + foreground: 'var(--code-foreground)', + }, sourcegraph: { blue: '#00CBEC', purple: '#A112FF',