From 6360fbb3b5dd9bed9d3ab8aee38d8d963ac48040 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Fri, 10 Jan 2025 23:29:38 -0800 Subject: [PATCH 1/3] feat(cody): improve context review logic Fixes a regression where the agent would mark a context as needed but got filtered out during the validation step - Simplify the context review logic by directly extracting context names from the response and fetching the full content for the requested files. - Remove unnecessary callbacks and optimize the context fetching process. - Add a helper function to check if the response is ready to answer. --- vscode/src/chat/agentic/DeepCody.ts | 72 +++++++++++++++-------------- 1 file changed, 38 insertions(+), 34 deletions(-) diff --git a/vscode/src/chat/agentic/DeepCody.ts b/vscode/src/chat/agentic/DeepCody.ts index 5fc7330d608c..ca695044d387 100644 --- a/vscode/src/chat/agentic/DeepCody.ts +++ b/vscode/src/chat/agentic/DeepCody.ts @@ -18,6 +18,7 @@ import { telemetryRecorder, wrapInActiveSpan, } from '@sourcegraph/cody-shared' +import { getContextFromRelativePath } from '../../commands/context/file-path' import { forkSignal } from '../../completions/utils' import { getCategorizedMentions, isUserAddedItem } from '../../prompt-builder/utils' import type { ChatBuilder } from '../chat-view/ChatBuilder' @@ -144,7 +145,6 @@ export class DeepCodyAgent { maxLoops = 2 ): Promise { span.setAttribute('sampled', true) - this.statusCallback?.onStart() const startTime = performance.now() await this.reviewLoop(requestID, span, chatAbortSignal, maxLoops) telemetryRecorder.recordEvent('cody.deep-cody.context', 'reviewed', { @@ -164,7 +164,6 @@ export class DeepCodyAgent { category: 'billable', }, }) - this.statusCallback?.onComplete() return this.context } @@ -179,7 +178,6 @@ export class DeepCodyAgent { this.stats.loop++ const newContext = await this.review(requestID, span, chatAbortSignal) if (!newContext.length) break - // Filter and add new context items in one pass const validItems = newContext.filter(c => c.title !== 'TOOLCONTEXT') this.context.push(...validItems) @@ -204,16 +202,14 @@ export class DeepCodyAgent { chatAbortSignal: AbortSignal ): Promise { const prompter = this.getPrompter(this.context) - const promptData = await prompter.makePrompt(this.chatBuilder, 1, this.promptMixins) + const { prompt } = await prompter.makePrompt(this.chatBuilder, 1, this.promptMixins) span.addEvent('sendReviewRequest') try { - const res = await this.processStream( - requestID, - promptData.prompt, - chatAbortSignal, - DeepCodyAgent.model - ) - if (!res) return [] + const res = await this.processStream(requestID, prompt, chatAbortSignal, DeepCodyAgent.model) + // If the response is empty or only contains the answer token, it's ready to answer. + if (!res || isReadyToAnswer(res)) { + return [] + } const results = await Promise.all( this.tools.map(async tool => { try { @@ -234,39 +230,43 @@ export class DeepCodyAgent { ) const reviewed = [] - - // Extract all the strings from between tags. - const valid = RawTextProcessor.extract(res, ACTIONS_TAGS.CONTEXT.toString()) - for (const contextName of valid || []) { - const foundValidatedItems = this.context.filter(c => c.uri.path.endsWith(contextName)) - for (const found of foundValidatedItems) { - reviewed.push({ ...found, source: ContextItemSource.Agentic }) + const currentContext = [ + ...this.context, + ...this.chatBuilder + .getDehydratedMessages() + .flatMap(m => (m.contextFiles ? [...m.contextFiles].reverse() : [])) + .filter(isDefined), + ] + // Extract context items that are enclosed with context tags from the response. + // We will validate the context items by checking if the context item is in the current context, + // which is a list of context that we have fetched in this round, and the ones from user's current + // chat session. + const contextNames = RawTextProcessor.extract(res, contextTag) + for (const contextName of contextNames) { + for (const item of currentContext) { + if (item.uri.path.endsWith(contextName)) { + // Try getting the full content for the requested file. + const file = (await getContextFromRelativePath(contextName)) || item + reviewed.push({ ...file, source: ContextItemSource.Agentic }) + } } } - - // Replace the current context list with the reviewed context. - if (valid.length + reviewed.length > 0) { - reviewed.push(...this.context.filter(c => isUserAddedItem(c))) + // When there are context items matched, we will replace the current context with + // the reviewed context list, but first we will make sure all the user added context + // items are nor removed from the updated context list. We will let the prompt builder + // at the final stage to do the unique context check. + if (reviewed.length > 0) { + const userAdded = this.context.filter(c => isUserAddedItem(c)) + reviewed.push(...userAdded) this.context = reviewed } - // If the response is empty or contains the known token, the context is sufficient. - if (res?.includes(ACTIONS_TAGS.ANSWER.toString())) { - // Process the response without generating any context items. - for (const tool of this.tools) { - tool.processResponse?.() - } - return reviewed - } - const newContextFetched = results.flat().filter(isDefined) this.stats.context = this.stats.context + newContextFetched.length return newContextFetched } catch (error) { await this.multiplexer.notifyTurnComplete() - logDebug('Deep Cody', `context review failed: ${error}`, { - verbose: { prompt: promptData.prompt, error }, - }) + logDebug('Deep Cody', `context review failed: ${error}`, { verbose: { prompt, error } }) return [] } } @@ -356,3 +356,7 @@ export class RawTextProcessor { return PromptString.join(prompts, connector) } } + +const answerTag = ACTIONS_TAGS.ANSWER.toString() +const contextTag = ACTIONS_TAGS.CONTEXT.toString() +const isReadyToAnswer = (text: string) => text === `<${answerTag}>` From 0f93919d07e527360965fe005270da4a26d19017 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Fri, 10 Jan 2025 23:52:41 -0800 Subject: [PATCH 2/3] add unit test --- vscode/src/chat/agentic/DeepCody.test.ts | 58 ++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/vscode/src/chat/agentic/DeepCody.test.ts b/vscode/src/chat/agentic/DeepCody.test.ts index b1107ee58cca..e254dd8623ff 100644 --- a/vscode/src/chat/agentic/DeepCody.test.ts +++ b/vscode/src/chat/agentic/DeepCody.test.ts @@ -212,4 +212,62 @@ describe('DeepCody', () => { expect(result.some(r => r.content === 'const example = "test";')).toBeFalsy() expect(result.some(r => r.content === 'const newExample = "test result";')).toBeTruthy() }) + + it('validates and preserves context items during review process', async () => { + // Mock a scenario where we have existing context items + const existingContext = [ + { + uri: URI.file('/path/to/file1.ts'), + type: 'file', + source: ContextItemSource.User, + content: 'const userAddedFile = "test";', + }, + { + uri: URI.file('/path/to/file2.ts'), + type: 'file', + source: ContextItemSource.Search, + content: 'const searchResult = "test";', + }, + ] satisfies ContextItem[] + + // Mock the chat messages to include context files + mockChatBuilder.getDehydratedMessages = vi.fn().mockReturnValue([ + { + speaker: 'human', + text: ps`test message`, + contextFiles: existingContext, + }, + ]) + + // Mock stream response that includes context validation tags + const mockStreamResponse = [ + { + type: 'change', + text: 'file1.tsfile2.tsnewfile.ts', + }, + { type: 'complete' }, + ] + + mockChatClient.chat = vi.fn().mockReturnValue(mockStreamResponse) + + // Create agent and run context retrieval + const agent = new DeepCodyAgent(mockChatBuilder, mockChatClient, mockStatusCallback) + const result = await agent.getContext( + 'deep-cody-test-validation-id', + new AbortController().signal, + existingContext + ) + + // Verify results + expect(mockChatClient.chat).toHaveBeenCalled() + + // Should preserve user-added context as is + expect(result.some(r => r.content === 'const userAddedFile = "test";')).toBeTruthy() + expect(result.some(r => r.source === ContextItemSource.User)).toBeTruthy() + + // Should include validated context from review + expect(result.some(r => r.content === 'const searchResult = "test";')).toBeTruthy() + // Should replace search context with agentic source during validation. + expect(result.filter(r => r.source === ContextItemSource.Search).length).toBe(0) + }) }) From 19dae1467c011d13a4fa2f1d768d59a6c3c7956b Mon Sep 17 00:00:00 2001 From: Beatrix <68532117+abeatrix@users.noreply.github.com> Date: Mon, 13 Jan 2025 07:38:23 -0800 Subject: [PATCH 3/3] Update DeepCody.ts Co-authored-by: Ara --- vscode/src/chat/agentic/DeepCody.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vscode/src/chat/agentic/DeepCody.ts b/vscode/src/chat/agentic/DeepCody.ts index ca695044d387..1cb238c88cac 100644 --- a/vscode/src/chat/agentic/DeepCody.ts +++ b/vscode/src/chat/agentic/DeepCody.ts @@ -253,7 +253,7 @@ export class DeepCodyAgent { } // When there are context items matched, we will replace the current context with // the reviewed context list, but first we will make sure all the user added context - // items are nor removed from the updated context list. We will let the prompt builder + // items are not removed from the updated context list. We will let the prompt builder // at the final stage to do the unique context check. if (reviewed.length > 0) { const userAdded = this.context.filter(c => isUserAddedItem(c))