Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(cody): fix chat context review logic #6602

Merged
merged 3 commits into from
Jan 13, 2025
Merged

fix(cody): fix chat context review logic #6602

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6360fbb
feat(cody): improve context review logic
abeatrix Jan 11, 2025
0f93919
add unit test
abeatrix Jan 11, 2025
19dae14
Update DeepCody.ts
abeatrix Jan 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 58 additions & 0 deletions vscode/src/chat/agentic/DeepCody.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: '<context_list>file1.ts</context_list><context_list>file2.ts</context_list><context_list>newfile.ts</context_list>',
},
{ 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)
})
})
72 changes: 38 additions & 34 deletions vscode/src/chat/agentic/DeepCody.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -144,7 +145,6 @@ export class DeepCodyAgent {
maxLoops = 2
): Promise<ContextItem[]> {
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', {
Expand All @@ -164,7 +164,6 @@ export class DeepCodyAgent {
category: 'billable',
},
})
this.statusCallback?.onComplete()
return this.context
}

Expand All @@ -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)
Expand All @@ -204,16 +202,14 @@ export class DeepCodyAgent {
chatAbortSignal: AbortSignal
): Promise<ContextItem[]> {
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 {
Expand All @@ -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 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))
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 []
}
}
Expand Down Expand Up @@ -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}>`
Loading