From 489ae3c307df0768698656a17ff2fb2aaccbf7b7 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Fri, 27 Dec 2024 16:26:08 -0800 Subject: [PATCH 1/5] fix(context-agent): add status callbacks back The setStatusCallback code was removed by #6469 by accident. This change is to add it back so the context agent (deep cody) can send its processing status back to webview 1. Added a `postMessageCallback` parameter to the `DeepCodyHandler` constructor to allow for posting empty messages during the context retrieval process. 2. Updated the `DeepCodyHandler` to set the status callback on the `DeepCodyAgent` instance, allowing for the posting of empty messages during the context retrieval process. 3. Updated the `AgentTools` interface to include the `postMessageCallback` parameter, ensuring that it is passed to the `DeepCodyHandler` constructor. 4. Updated the `ChatController` to pass the `postEmptyMessageInProgress` callback to the `DeepCodyHandler` constructor, enabling the posting of empty messages during the context retrieval process. --- vscode/src/chat/chat-view/ChatController.ts | 1 + .../src/chat/chat-view/handlers/DeepCodyHandler.ts | 8 +++++--- vscode/src/chat/chat-view/handlers/interfaces.ts | 1 + vscode/src/chat/chat-view/handlers/registry.ts | 14 ++++++++++++-- 4 files changed, 19 insertions(+), 5 deletions(-) diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts index de853295ecbd..01511bbabc75 100644 --- a/vscode/src/chat/chat-view/ChatController.ts +++ b/vscode/src/chat/chat-view/ChatController.ts @@ -807,6 +807,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv editor: this.editor, chatClient: this.chatClient, codyToolProvider: this.toolProvider, + postMessageCallback: model => this.postEmptyMessageInProgress(model), }) recorder.setIntentInfo({ diff --git a/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts b/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts index ea7b3108723f..47bdc8483f8d 100644 --- a/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts +++ b/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts @@ -21,7 +21,8 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler { contextRetriever: Pick, editor: ChatControllerOptions['editor'], chatClient: ChatControllerOptions['chatClient'], - private toolProvider: CodyToolProvider + private toolProvider: CodyToolProvider, + private postMessageCallback: (model: string) => void ) { super(modelId, contextRetriever, editor, chatClient) } @@ -66,13 +67,14 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler { } const baseContext = baseContextResult.contextItems ?? [] - const codyAgent = new DeepCodyAgent( + const agent = new DeepCodyAgent( chatBuilder, this.chatClient, this.toolProvider.getTools(), baseContext ) - const agenticContext = await codyAgent.getContext(requestID, signal) + agent.setStatusCallback(model => this.postMessageCallback(model)) + const agenticContext = await agent.getContext(requestID, signal) return { contextItems: [...baseContext, ...agenticContext] } } } diff --git a/vscode/src/chat/chat-view/handlers/interfaces.ts b/vscode/src/chat/chat-view/handlers/interfaces.ts index eb797a45db21..7f44c6f4815d 100644 --- a/vscode/src/chat/chat-view/handlers/interfaces.ts +++ b/vscode/src/chat/chat-view/handlers/interfaces.ts @@ -17,6 +17,7 @@ export interface AgentTools { editor: ChatControllerOptions['editor'] chatClient: ChatControllerOptions['chatClient'] codyToolProvider: CodyToolProvider + postMessageCallback: (model: string) => void } /** diff --git a/vscode/src/chat/chat-view/handlers/registry.ts b/vscode/src/chat/chat-view/handlers/registry.ts index a9ce38348aad..4c507b2df6af 100644 --- a/vscode/src/chat/chat-view/handlers/registry.ts +++ b/vscode/src/chat/chat-view/handlers/registry.ts @@ -25,8 +25,18 @@ export function getAgent(id: string, tools: AgentTools): AgentHandler { registerAgent( 'sourcegraph::2023-06-01::deep-cody', - (id: string, { contextRetriever, editor, chatClient, codyToolProvider }: AgentTools) => - new DeepCodyHandler(id, contextRetriever, editor, chatClient, codyToolProvider) + ( + id: string, + { contextRetriever, editor, chatClient, codyToolProvider, postMessageCallback }: AgentTools + ) => + new DeepCodyHandler( + id, + contextRetriever, + editor, + chatClient, + codyToolProvider, + postMessageCallback + ) ) registerAgent('search', (_id: string, _tools: AgentTools) => new SearchHandler()) registerAgent( From fe26386a7d74654465ae4ef18f2e99dcd28952d8 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Fri, 27 Dec 2024 17:32:02 -0800 Subject: [PATCH 2/5] pass delegate to computeContext --- vscode/src/chat/chat-view/ChatController.ts | 1 - vscode/src/chat/chat-view/handlers/ChatHandler.ts | 2 ++ .../src/chat/chat-view/handlers/DeepCodyHandler.ts | 12 ++++++++---- vscode/src/chat/chat-view/handlers/interfaces.ts | 1 - vscode/src/chat/chat-view/handlers/registry.ts | 14 ++------------ 5 files changed, 12 insertions(+), 18 deletions(-) diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts index 01511bbabc75..de853295ecbd 100644 --- a/vscode/src/chat/chat-view/ChatController.ts +++ b/vscode/src/chat/chat-view/ChatController.ts @@ -807,7 +807,6 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv editor: this.editor, chatClient: this.chatClient, codyToolProvider: this.toolProvider, - postMessageCallback: model => this.postEmptyMessageInProgress(model), }) recorder.setIntentInfo({ diff --git a/vscode/src/chat/chat-view/handlers/ChatHandler.ts b/vscode/src/chat/chat-view/handlers/ChatHandler.ts index 6edbc146c15b..839f074ed22d 100644 --- a/vscode/src/chat/chat-view/handlers/ChatHandler.ts +++ b/vscode/src/chat/chat-view/handlers/ChatHandler.ts @@ -57,6 +57,7 @@ export class ChatHandler implements AgentHandler { { text: inputText, mentions }, editorState, chatBuilder, + delegate, signal ) if (contextResult.error) { @@ -223,6 +224,7 @@ export class ChatHandler implements AgentHandler { { text, mentions }: HumanInput, editorState: SerializedPromptEditorState | null, _chatBuilder: ChatBuilder, + _delegate: AgentHandlerDelegate, signal?: AbortSignal ): Promise<{ contextItems?: ContextItem[] diff --git a/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts b/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts index 47bdc8483f8d..a05de8d18b2e 100644 --- a/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts +++ b/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts @@ -13,7 +13,7 @@ import type { ChatControllerOptions } from '../ChatController' import type { ContextRetriever } from '../ContextRetriever' import type { HumanInput } from '../context' import { ChatHandler } from './ChatHandler' -import type { AgentHandler } from './interfaces' +import type { AgentHandler, AgentHandlerDelegate } from './interfaces' export class DeepCodyHandler extends ChatHandler implements AgentHandler { constructor( @@ -21,8 +21,7 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler { contextRetriever: Pick, editor: ChatControllerOptions['editor'], chatClient: ChatControllerOptions['chatClient'], - private toolProvider: CodyToolProvider, - private postMessageCallback: (model: string) => void + private toolProvider: CodyToolProvider ) { super(modelId, contextRetriever, editor, chatClient) } @@ -39,6 +38,7 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler { { text, mentions }: HumanInput, editorState: SerializedPromptEditorState | null, chatBuilder: ChatBuilder, + delegate: AgentHandlerDelegate, signal: AbortSignal ): Promise<{ contextItems?: ContextItem[] @@ -50,6 +50,7 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler { { text, mentions }, editorState, chatBuilder, + delegate, signal ) const isEnabled = chatBuilder.getMessages().length < 4 @@ -73,7 +74,10 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler { this.toolProvider.getTools(), baseContext ) - agent.setStatusCallback(model => this.postMessageCallback(model)) + // Use delegate instead of callback + agent.setStatusCallback(model => { + delegate.postMessageInProgress({ speaker: 'assistant', model }) + }) const agenticContext = await agent.getContext(requestID, signal) return { contextItems: [...baseContext, ...agenticContext] } } diff --git a/vscode/src/chat/chat-view/handlers/interfaces.ts b/vscode/src/chat/chat-view/handlers/interfaces.ts index 7f44c6f4815d..eb797a45db21 100644 --- a/vscode/src/chat/chat-view/handlers/interfaces.ts +++ b/vscode/src/chat/chat-view/handlers/interfaces.ts @@ -17,7 +17,6 @@ export interface AgentTools { editor: ChatControllerOptions['editor'] chatClient: ChatControllerOptions['chatClient'] codyToolProvider: CodyToolProvider - postMessageCallback: (model: string) => void } /** diff --git a/vscode/src/chat/chat-view/handlers/registry.ts b/vscode/src/chat/chat-view/handlers/registry.ts index 4c507b2df6af..a9ce38348aad 100644 --- a/vscode/src/chat/chat-view/handlers/registry.ts +++ b/vscode/src/chat/chat-view/handlers/registry.ts @@ -25,18 +25,8 @@ export function getAgent(id: string, tools: AgentTools): AgentHandler { registerAgent( 'sourcegraph::2023-06-01::deep-cody', - ( - id: string, - { contextRetriever, editor, chatClient, codyToolProvider, postMessageCallback }: AgentTools - ) => - new DeepCodyHandler( - id, - contextRetriever, - editor, - chatClient, - codyToolProvider, - postMessageCallback - ) + (id: string, { contextRetriever, editor, chatClient, codyToolProvider }: AgentTools) => + new DeepCodyHandler(id, contextRetriever, editor, chatClient, codyToolProvider) ) registerAgent('search', (_id: string, _tools: AgentTools) => new SearchHandler()) registerAgent( From fcbc98e564e991e6b22bd4464a69fb8cd9755488 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Fri, 27 Dec 2024 18:08:56 -0800 Subject: [PATCH 3/5] Display context before sending the request --- vscode/src/chat/chat-view/handlers/ChatHandler.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vscode/src/chat/chat-view/handlers/ChatHandler.ts b/vscode/src/chat/chat-view/handlers/ChatHandler.ts index 839f074ed22d..e45b190ad859 100644 --- a/vscode/src/chat/chat-view/handlers/ChatHandler.ts +++ b/vscode/src/chat/chat-view/handlers/ChatHandler.ts @@ -82,6 +82,8 @@ export class ChatHandler implements AgentHandler { recorder.recordChatQuestionExecuted(corpusContext, { addMetadata: true, current: span }) signal.throwIfAborted() + // Display context in webview before sending the request. + delegate.postMessageInProgress({ speaker: 'assistant', model: this.modelId }) this.streamAssistantResponse(requestID, prompt, this.modelId, signal, chatBuilder, delegate) } From e13a5f5d3c22803598ed328c0eb41299cb986904 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Fri, 27 Dec 2024 18:18:29 -0800 Subject: [PATCH 4/5] revert --- vscode/src/chat/chat-view/handlers/ChatHandler.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/vscode/src/chat/chat-view/handlers/ChatHandler.ts b/vscode/src/chat/chat-view/handlers/ChatHandler.ts index e45b190ad859..839f074ed22d 100644 --- a/vscode/src/chat/chat-view/handlers/ChatHandler.ts +++ b/vscode/src/chat/chat-view/handlers/ChatHandler.ts @@ -82,8 +82,6 @@ export class ChatHandler implements AgentHandler { recorder.recordChatQuestionExecuted(corpusContext, { addMetadata: true, current: span }) signal.throwIfAborted() - // Display context in webview before sending the request. - delegate.postMessageInProgress({ speaker: 'assistant', model: this.modelId }) this.streamAssistantResponse(requestID, prompt, this.modelId, signal, chatBuilder, delegate) } From 697868f9553417f82904ea8167cfc388fa3c175d Mon Sep 17 00:00:00 2001 From: Beyang Liu Date: Mon, 30 Dec 2024 13:57:24 -0800 Subject: [PATCH 5/5] add postStatuses delegate method and use it to convey status update --- vscode/src/chat/agentic/CodyChatAgent.ts | 44 ++++++++----------- vscode/src/chat/chat-view/ChatController.ts | 4 ++ .../chat-view/handlers/DeepCodyHandler.ts | 8 ++-- .../src/chat/chat-view/handlers/interfaces.ts | 2 + 4 files changed, 28 insertions(+), 30 deletions(-) diff --git a/vscode/src/chat/agentic/CodyChatAgent.ts b/vscode/src/chat/agentic/CodyChatAgent.ts index 03e71241cf39..a41afc8be1b5 100644 --- a/vscode/src/chat/agentic/CodyChatAgent.ts +++ b/vscode/src/chat/agentic/CodyChatAgent.ts @@ -3,6 +3,7 @@ import { type ChatClient, type ContextItem, type Message, + type ProcessingStep, type PromptMixin, type PromptString, newPromptMixin, @@ -19,18 +20,34 @@ export abstract class CodyChatAgent { protected readonly promptMixins: PromptMixin[] = [] protected readonly toolHandlers: Map protected statusCallback?: ToolStatusCallback - protected postMessageCallback?: (model: string) => void + private stepsManager?: ProcessManager constructor( protected readonly chatBuilder: ChatBuilder, protected readonly chatClient: Pick, protected readonly tools: CodyTool[], - protected context: ContextItem[] = [] + protected context: ContextItem[], + statusUpdateCallback?: (steps: ProcessingStep[]) => void ) { // Initialize handlers and mixins in constructor this.toolHandlers = new Map(tools.map(tool => [tool.config.tags.tag.toString(), tool])) this.initializeMultiplexer() this.promptMixins.push(newPromptMixin(this.buildPrompt())) + + this.stepsManager = new ProcessManager(steps => { + statusUpdateCallback?.(steps) + }) + this.statusCallback = { + onStart: () => { + this.stepsManager?.initializeStep() + }, + onStream: (toolName, content) => { + this.stepsManager?.addStep(toolName, content) + }, + onComplete: (toolName, error) => { + this.stepsManager?.completeStep(toolName, error) + }, + } } protected initializeMultiplexer(): void { @@ -86,29 +103,6 @@ export abstract class CodyChatAgent { return new DefaultPrompter(explicitMentions, implicitMentions.slice(-MAX_SEARCH_ITEMS)) } - public setStatusCallback(postMessage: (model: string) => void): void { - this.postMessageCallback = postMessage - const model = this.chatBuilder.selectedModel ?? '' - - // Create a steps manager to handle state updates efficiently - const stepsManager = new ProcessManager(steps => { - this.chatBuilder.setLastMessageProcesses(steps) - this.postMessageCallback?.(model) - }) - - this.statusCallback = { - onStart: () => { - stepsManager.initializeStep() - }, - onStream: (toolName, content) => { - stepsManager.addStep(toolName, content) - }, - onComplete: (toolName, error) => { - stepsManager.completeStep(toolName, error) - }, - } - } - // Abstract methods that must be implemented by derived classes protected abstract buildPrompt(): PromptString } diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts index de853295ecbd..b42d1bcf3e95 100644 --- a/vscode/src/chat/chat-view/ChatController.ts +++ b/vscode/src/chat/chat-view/ChatController.ts @@ -18,6 +18,7 @@ import { type Guardrails, ModelUsage, type NLSSearchDynamicFilter, + type ProcessingStep, PromptString, type SerializedChatInteraction, type SerializedChatTranscript, @@ -842,6 +843,9 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv messageInProgress = message this.postViewTranscript(message) }, + postStatuses: (steps: ProcessingStep[]): void => { + this.chatBuilder.setLastMessageProcesses(steps) + }, postDone: (op?: { abort: boolean }): void => { if (op?.abort) { this.handleAbort() diff --git a/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts b/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts index a05de8d18b2e..16b9e346cd5f 100644 --- a/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts +++ b/vscode/src/chat/chat-view/handlers/DeepCodyHandler.ts @@ -1,6 +1,7 @@ import { type ContextItem, FeatureFlag, + type ProcessingStep, type SerializedPromptEditorState, featureFlagProvider, storeLastValue, @@ -72,12 +73,9 @@ export class DeepCodyHandler extends ChatHandler implements AgentHandler { chatBuilder, this.chatClient, this.toolProvider.getTools(), - baseContext + baseContext, + (steps: ProcessingStep[]) => delegate.postStatuses(steps) ) - // Use delegate instead of callback - agent.setStatusCallback(model => { - delegate.postMessageInProgress({ speaker: 'assistant', model }) - }) const agenticContext = await agent.getContext(requestID, signal) return { contextItems: [...baseContext, ...agenticContext] } } diff --git a/vscode/src/chat/chat-view/handlers/interfaces.ts b/vscode/src/chat/chat-view/handlers/interfaces.ts index eb797a45db21..112f160d364d 100644 --- a/vscode/src/chat/chat-view/handlers/interfaces.ts +++ b/vscode/src/chat/chat-view/handlers/interfaces.ts @@ -2,6 +2,7 @@ import type { Span } from '@opentelemetry/api' import type { ChatMessage, ContextItem, + ProcessingStep, PromptString, SerializedPromptEditorState, } from '@sourcegraph/cody-shared' @@ -24,6 +25,7 @@ export interface AgentTools { */ export interface AgentHandlerDelegate { postError(error: Error, type?: MessageErrorType): void + postStatuses(steps: ProcessingStep[]): void postMessageInProgress(message: ChatMessage): void postDone(ops?: { abort: boolean }): void }