Inline Fix: Improve response quality (#376)

agent/src/editor.ts

@@ -2,6 +2,7 @@ import { URI } from 'vscode-uri'
 
 import {
     ActiveTextEditor,
+    ActiveTextEditorDiagnostic,
     ActiveTextEditorSelection,
     ActiveTextEditorViewControllers,
     ActiveTextEditorVisibleContent,
@@ -78,6 +79,10 @@ export class AgentEditor implements Editor {
         return this.getActiveTextEditorSelection()
     }
 
+    public getActiveTextEditorDiagnosticsForRange(): ActiveTextEditorDiagnostic[] | null {
+        throw new Error('Method not implemented.')
+    }
+
     public getActiveTextEditorVisibleContent(): ActiveTextEditorVisibleContent | null {
         const document = this.activeDocument()
         if (document === undefined) {

cli/src/client/index.ts

@@ -49,14 +49,14 @@ export async function getClient({ codebase, endpoint, context: contextType, debu
         process.exit(1)
     }
 
-    const intentDetector = new SourcegraphIntentDetectorClient(sourcegraphClient)
-
     const completionsClient = new SourcegraphNodeCompletionsClient({
         serverEndpoint: endpoint,
         accessToken,
         debugEnable: debug,
         customHeaders: {},
     })
 
+    const intentDetector = new SourcegraphIntentDetectorClient(sourcegraphClient, completionsClient)
+
     return { codebaseContext, intentDetector, completionsClient }
 }

lib/shared/src/chat/client.ts

@@ -89,7 +89,7 @@ export async function createClient({
 
         const codebaseContext = new CodebaseContext(config, config.codebase, embeddingsSearch, null, null)
 
-        const intentDetector = new SourcegraphIntentDetectorClient(graphqlClient)
+        const intentDetector = new SourcegraphIntentDetectorClient(graphqlClient, completionsClient)
 
         const transcript = initialTranscript || new Transcript()
 

lib/shared/src/chat/recipes/fixup.ts

@@ -1,13 +1,54 @@
-import { CodebaseContext } from '../../codebase-context'
-import { ContextMessage } from '../../codebase-context/messages'
+import { ContextMessage, getContextMessageWithResponse } from '../../codebase-context/messages'
+import { ActiveTextEditorSelection } from '../../editor'
+import { IntentClassificationOption } from '../../intent-detector'
 import { MAX_CURRENT_FILE_TOKENS, MAX_HUMAN_INPUT_TOKENS } from '../../prompt/constants'
+import { populateCodeContextTemplate, populateCurrentEditorDiagnosticsTemplate } from '../../prompt/templates'
 import { truncateText, truncateTextStart } from '../../prompt/truncation'
 import { BufferedBotResponseSubscriber } from '../bot-response-multiplexer'
 import { Interaction } from '../transcript/interaction'
 
-import { contentSanitizer, numResults } from './helpers'
+import { contentSanitizer, getContextMessagesFromSelection } from './helpers'
 import { Recipe, RecipeContext, RecipeID } from './recipe'
 
+type FixupIntent = 'add' | 'edit' | 'delete' | 'fix' | 'test' | 'document'
+const FixupIntentClassification: IntentClassificationOption<FixupIntent>[] = [
+    {
+        id: 'add',
+        description: 'Add to the selected code',
+        examplePrompts: ['Add a function that concatonates two strings', 'Add error handling'],
+    },
+    {
+        id: 'edit',
+        description: 'Edit part of the selected code',
+        examplePrompts: ['Edit this code', 'Change this code', 'Update this code'],
+    },
+    {
+        id: 'delete',
+        description: 'Delete a part of the selection code',
+        examplePrompts: ['Delete these comments', 'Remove log statements'],
+    },
+    {
+        id: 'fix',
+        description: 'Fix a problem in a part of the selected code',
+        examplePrompts: ['Implement this TODO', 'Fix this code'],
+    },
+    {
+        id: 'document',
+        description: 'Generate documentation for parts of the selected code.',
+        examplePrompts: ['Add a docstring for this function', 'Write comments to explain this code'],
+    },
+]
+
+const PromptIntentInstruction: Record<FixupIntent, string> = {
+    add: 'The user wants you to add to the selected code by following their instructions.',
+    edit: 'The user wants you to replace parts of the selected code by following their instructions.',
+    delete: 'The user wants you to remove parts of the selected code by following their instructions.',
+    fix: 'The user wants you to correct a problem in the selected code by following their instructions.',
+    document:
+        'The user wants you to add documentation or comments to the selected code by following their instructions.',
+    test: 'The user wants you to add, update or fix a test by following their instructions',
+}
+
 export class Fixup implements Recipe {
     public id: RecipeID = 'fixup'
 
@@ -27,15 +68,22 @@ export class Fixup implements Recipe {
             return null
         }
 
-        // Reconstruct Cody's prompt using user's context
+        const intent = await this.getIntent(humanChatInput, context)
+
+        // It is possible to trigger this recipe from the sidebar without any input.
+        // TODO: Consider deprecating this functionality once inline fixups and non-stop fixups and consolidated.
+        const promptInstruction =
+            humanChatInput.length > 0
+                ? truncateText(humanChatInput, MAX_HUMAN_INPUT_TOKENS)
+                : "You should infer your instructions from the users' selection"
+
+        // Reconstruct Cody's prompt using user's context and intent
         // Replace placeholders in reverse order to avoid collisions if a placeholder occurs in the input
         // TODO: Move prompt suffix from recipe to chat view. It has other subscribers.
         const promptText = Fixup.prompt
-            .replace('{humanInput}', truncateText(humanChatInput, MAX_HUMAN_INPUT_TOKENS))
-            .replace('{responseMultiplexerPrompt}', context.responseMultiplexer.prompt())
-            .replace('{truncateFollowingText}', truncateText(selection.followingText, quarterFileContext))
+            .replace('{humanInput}', promptInstruction)
+            .replace('{intent}', PromptIntentInstruction[intent])
             .replace('{selectedText}', selection.selectedText)
-            .replace('{truncateTextStart}', truncateTextStart(selection.precedingText, quarterFileContext))
             .replace('{fileName}', selection.fileName)
 
         context.responseMultiplexer.sub(
@@ -67,36 +115,133 @@ export class Fixup implements Recipe {
                     speaker: 'assistant',
                     prefix: 'Check your document for updates from Cody.\n',
                 },
-                this.getContextMessages(selection.selectedText, context.codebaseContext),
+                this.getContextFromIntent(intent, selection, quarterFileContext, context),
                 []
             )
         )
     }
 
-    // Get context from editor
-    private async getContextMessages(text: string, codebaseContext: CodebaseContext): Promise<ContextMessage[]> {
-        const contextMessages: ContextMessage[] = await codebaseContext.getContextMessages(text, numResults)
-        return contextMessages
+    private async getIntent(humanChatInput: string, context: RecipeContext): Promise<FixupIntent> {
+        /**
+         * TODO(umpox): We should probably find a shorter way of detecting intent when possible.
+         * Possible methods:
+         * - Input -> Match first word against update|fix|add|delete verbs
+         * - Context -> Infer intent from context, e.g. Current file is a test -> Test intent, Current selection is a comment symbol -> Documentation intent
+         */
+        const intent = await context.intentDetector.classifyIntentFromOptions(
+            humanChatInput,
+            FixupIntentClassification,
+            'edit'
+        )
+        return intent
+    }
+
+    private async getContextFromIntent(
+        intent: FixupIntent,
+        selection: ActiveTextEditorSelection,
+        quarterFileContext: number,
+        context: RecipeContext
+    ): Promise<ContextMessage[]> {
+        const truncatedPrecedingText = truncateTextStart(selection.precedingText, quarterFileContext)
+        const truncatedFollowingText = truncateText(selection.followingText, quarterFileContext)
+
+        // Disable no case declarations because we get better type checking with a switch case
+        /* eslint-disable no-case-declarations */
+        switch (intent) {
+            /**
+             * Intents that are focused on producing new code.
+             * They have a broad set of possible instructions, so we fetch a broad amount of code context files.
+             * Non-code files are not considered as including Markdown syntax seems to lead to more hallucinations and poorer output quality.
+             *
+             * TODO(umpox): We fetch similar context for both cases here
+             * We should investigate how we can improve each individual case.
+             * Are these fundamentally the same? Is the primary benefit here that we can provide more specific instructions to Cody?
+             */
+            case 'add':
+            case 'edit':
+                return getContextMessagesFromSelection(
+                    selection.selectedText,
+                    truncatedPrecedingText,
+                    truncatedFollowingText,
+                    selection,
+                    context.codebaseContext
+                )
+            /**
+             * The fix intent is similar to adding or editing code, but with additional context that we can include from the editor.
+             */
+            case 'fix':
+                const range =
+                    context.editor.getActiveTextEditor()?.selectionRange ||
+                    context.editor.controllers?.inline?.selectionRange
+                const diagnostics = range ? context.editor.getActiveTextEditorDiagnosticsForRange(range) || [] : []
+                const errorsAndWarnings = diagnostics.filter(({ type }) => type === 'error' || type === 'warning')
+
+                return getContextMessagesFromSelection(
+                    selection.selectedText,
+                    truncatedPrecedingText,
+                    truncatedFollowingText,
+                    selection,
+                    context.codebaseContext
+                ).then(messages =>
+                    messages.concat(
+                        errorsAndWarnings.flatMap(diagnostic =>
+                            getContextMessageWithResponse(
+                                populateCurrentEditorDiagnosticsTemplate(diagnostic, selection.fileName),
+                                selection
+                            )
+                        )
+                    )
+                )
+            /**
+             * The test intent is unique in that we likely want to be much more specific in that context that we retrieve.
+             * TODO(umpox): How can infer the current testing dependencies, etc?
+             */
+            case 'test':
+                // Currently the same as add|edit|fix
+                return getContextMessagesFromSelection(
+                    selection.selectedText,
+                    truncatedPrecedingText,
+                    truncatedFollowingText,
+                    selection,
+                    context.codebaseContext
+                )
+            /**
+             * Intents that are focused primarily on updating code within the current file and selection.
+             * Providing a much more focused context window here seems to provide better quality responses.
+             */
+            case 'delete':
+            case 'document':
+                return Promise.resolve(
+                    [truncatedPrecedingText, truncatedFollowingText].flatMap(text =>
+                        getContextMessageWithResponse(
+                            populateCodeContextTemplate(text, selection.fileName, selection.repoName),
+                            selection
+                        )
+                    )
+                )
+        }
+        /* eslint-enable no-case-declarations */
     }
 
     // Prompt Templates
     public static readonly prompt = `
-    This is part of the file {fileName}. The part of the file I have selected is highlighted with <selection> tags. You are helping me to work on that part as my coding assistant.
-    Follow the instructions in the selected part plus the additional instructions to produce a rewritten replacement for only the selected part.
-    Put the rewritten replacement inside <selection> tags. I only want to see the code within <selection>.
-    Do not move code from outside the selection into the selection in your reply.
-    Do not remove code inside the <selection> tags that might be being used by the code outside the <selection> tags.
-    It is OK to provide some commentary within the replacement <selection>.
-    It is not acceptable to enclose the rewritten replacement with markdowns.
-    Only provide me with the replacement <selection> and nothing else.
-    If it doesn't make sense, you do not need to provide <selection>. Instead, tell me how I can help you to understand my request.
-
-    \`\`\`
-    {truncateTextStart}<selection>{selectedText}</selection>{truncateFollowingText}
-    \`\`\`
-
-    Additional Instruction:
-    - {humanInput}
-    - {responseMultiplexerPrompt}
-`
+- You are an AI programming assistant who is an expert in updating code to meet given instructions.
+- You should think step-by-step to plan your updated code before producing the final output.
+- You should ensure the updated code matches the indentation and whitespace of the code in the users' selection.
+- Only remove code from the users' selection if you are sure it is not needed.
+- It is not acceptable to use Markdown in your response. You should not produce Markdown-formatted code blocks.
+- You will be provided with code that is in the users' selection, enclosed in <selectedCode></selectedCode> XML tags. You must use this code to help you plan your updated code.
+- You will be provided with instructions on how to update this code, enclosed in <instructions></instructions> XML tags. You must follow these instructions carefully and to the letter.
+- Enclose your response in <selection></selection> XML tags. Do not provide anything else.
+
+This is part of the file {fileName}.
+
+The user has the following code in their selection:
+<selectedCode>{selectedText}</selectedCode>
+
+{intent}
+Provide your generated code using the following instructions:
+<instructions>
+{humanInput}
+</instructions>`
 }

lib/shared/src/chat/useClient.ts

@@ -134,7 +134,7 @@ export const useClient = ({
         const completionsClient = new SourcegraphBrowserCompletionsClient(config)
         const chatClient = new ChatClient(completionsClient)
         const graphqlClient = new SourcegraphGraphQLAPIClient(config)
-        const intentDetector = new SourcegraphIntentDetectorClient(graphqlClient)
+        const intentDetector = new SourcegraphIntentDetectorClient(graphqlClient, completionsClient)
 
         return { graphqlClient, chatClient, intentDetector }
     }, [config])

lib/shared/src/editor/index.ts

@@ -28,6 +28,15 @@ export interface ActiveTextEditorSelection {
     followingText: string
 }
 
+export type ActiveTextEditorDiagnosticType = 'error' | 'warning' | 'information' | 'hint'
+
+export interface ActiveTextEditorDiagnostic {
+    type: ActiveTextEditorDiagnosticType
+    range: ActiveTextEditorSelectionRange
+    text: string
+    message: string
+}
+
 export interface ActiveTextEditorVisibleContent {
     content: string
     fileName: string
@@ -37,6 +46,7 @@ export interface ActiveTextEditorVisibleContent {
 
 export interface VsCodeInlineController {
     selection: ActiveTextEditorSelection | null
+    selectionRange: ActiveTextEditorSelectionRange | null
     error(): Promise<void>
 }
 
@@ -92,6 +102,10 @@ export interface Editor<
      * Gets the active text editor's selection, or the entire file if the selected range is empty.
      */
     getActiveTextEditorSelectionOrEntireFile(): ActiveTextEditorSelection | null
+    /**
+     * Get diagnostics (errors, warnings, hints) for a range within the active text editor.
+     */
+    getActiveTextEditorDiagnosticsForRange(range: ActiveTextEditorSelectionRange): ActiveTextEditorDiagnostic[] | null
 
     getActiveTextEditorVisibleContent(): ActiveTextEditorVisibleContent | null
     replaceSelection(fileName: string, selectedText: string, replacement: string): Promise<void>
@@ -129,6 +143,10 @@ export class NoopEditor implements Editor {
         return null
     }
 
+    public getActiveTextEditorDiagnosticsForRange(): ActiveTextEditorDiagnostic[] | null {
+        return null
+    }
+
     public getActiveTextEditorVisibleContent(): ActiveTextEditorVisibleContent | null {
         return null
     }