feat: add context file via @ in chat input (#1631)

Continue from #1549 Close #1523 https://github.com/sourcegraph/cody/assets/68532117/0d2d868a-8031-4067-b933-705f331f539f This adds PR support for passing context files to recipe from chat input. The context files are extracted from the codebase and passed to the recipe context. - Create context messages from context files - Generate display text including context file names - Add getTextEditorContentForContextFile method in editor to get content for context files - Update editor interface to support getting content for context files - Add ContextFile type to represent context files - Add userInputContextFiles property to RecipeContext for passing context files - Update RecipeContext to accept ContextFile[] - Generate context messages from ContextFile[] - Update chat question recipe to support getContextFilesContext - Pass editor context as ContextFile[] to custom prompt and chat question handler > This PR focuses on implementing the feature that allows users to add local workspace file via @ command via chat box. ## Next - [ ] A widget that allows user to toggle between including enhanced context with chat questions or not #1524 - [ ] Replace current context file display widget with the new UI for enhanced context #1525 ## Test plan  - [ ] In the input box inside the new Chat view, add files from the current workspace using the `@` command - [ ] Typing `@` will display a pop up where you can select files in the current workspace - [ ] Typing `@` without additional character will show a list of currently opened files for you to choose - [ ] Typing `@` follow by additional character will update the file list and display results that matches the file path and symbol names when available - [ ] Allow you to attach multiple files via @ - [ ] Ask a question with a file attached using the `@` command - [ ] Cody should have context about the file or symbol you attached using the `@` command - [ ] The @-file is a link that's clickable in user's display text ![image](https://github.com/sourcegraph/cody/assets/68532117/97459faf-2e64-4a08-954a-b4d747c5fa69) --------- Co-authored-by: Tom Ross <tom@umpox.com>
sourcegraph · Nov 10, 2023 · 0e0574f · 0e0574f
1 parent ec008a8
commit 0e0574f
Show file tree

Hide file tree

Showing 29 changed files with 1,056 additions and 151 deletions.
diff --git a/agent/src/editor.ts b/agent/src/editor.ts
@@ -1,4 +1,5 @@
 import type * as vscode from 'vscode'
+import { URI } from 'vscode-uri'
 
 import {
     ActiveTextEditor,
@@ -50,6 +51,15 @@ export class AgentEditor implements Editor {
         }
     }
 
+    public async getTextEditorContentForFile(uri: URI): Promise<string | undefined> {
+        if (!uri) {
+            return Promise.resolve(undefined)
+        }
+
+        const doc = this.agent.workspace.getDocument(uri.fsPath)
+        return Promise.resolve(doc?.content)
+    }
+
     public getActiveTextEditorSelection(): ActiveTextEditorSelection | null {
         const document = this.activeDocument()
         if (document?.content === undefined || document.selection === undefined) {

diff --git a/lib/shared/src/chat/prompts/display-text.test.ts b/lib/shared/src/chat/prompts/display-text.test.ts
@@ -0,0 +1,81 @@
+import { describe, expect, it } from 'vitest'
+
+import { replaceFileNameWithMarkdownLink } from './display-text'
+
+describe('replaceFileNameWithMarkdownLink', () => {
+    it('replaces file name with markdown link', () => {
+        const text = 'Hello @test.js'
+
+        const result = replaceFileNameWithMarkdownLink(text, '@test.js', '/path/to/test.js')
+
+        expect(result).toEqual('Hello [_@test.js_](vscode://file/path/to/test.js)')
+    })
+
+    it('respects spaces in file name', () => {
+        const text = 'Loaded @my file.js'
+
+        const result = replaceFileNameWithMarkdownLink(text, '@my file.js', '/path/to/my file.js')
+
+        expect(result).toEqual('Loaded [_@my file.js_](vscode://file/path/to/my file.js)')
+    })
+
+    it('returns original text if no match', () => {
+        const text = 'No file name'
+
+        const result = replaceFileNameWithMarkdownLink(text, '@test.js', '/path/to/test.js')
+
+        expect(result).toEqual(text)
+    })
+
+    it('handles special characters in path', () => {
+        const text = 'Loaded @test.js'
+
+        const result = replaceFileNameWithMarkdownLink(text, '@test.js', '/path/with/@#special$chars.js')
+
+        expect(result).toEqual('Loaded [_@test.js_](vscode://file/path/with/@#special$chars.js)')
+    })
+
+    it('handles line numbers', () => {
+        const text = 'Error in @test.js'
+
+        const result = replaceFileNameWithMarkdownLink(text, '@test.js', '/path/test.js', 10)
+
+        expect(result).toEqual('Error in [_@test.js_](vscode://file/path/test.js:10)')
+    })
+
+    it('handles names that showed up more than once', () => {
+        const text = 'Compare and explain @foo.js and @bar.js. What does @foo.js do?'
+
+        const result = replaceFileNameWithMarkdownLink(text, '@foo.js', '/path/foo.js', 10)
+
+        expect(result).toEqual(
+            'Compare and explain [_@foo.js_](vscode://file/path/foo.js:10) and @bar.js. What does [_@foo.js_](vscode://file/path/foo.js:10) do?'
+        )
+    })
+
+    it('ignore repeated file names that are followed by another character', () => {
+        const text = 'Compare and explain @foo.js and @bar.js. What does @foo.js#FooBar() do?'
+
+        const result = replaceFileNameWithMarkdownLink(text, '@foo.js', '/path/foo.js', 10)
+
+        expect(result).toEqual(
+            'Compare and explain [_@foo.js_](vscode://file/path/foo.js:10) and @bar.js. What does @foo.js#FooBar() do?'
+        )
+    })
+
+    // FAILING - NEED TO BE FIXED
+    it('handles file names with line number and symbol name', () => {
+        const text = '@vscode/src/logged-rerank.ts:7-23#getRerankWithLog() what does this do'
+
+        const result = replaceFileNameWithMarkdownLink(
+            text,
+            '@vscode/src/logged-rerank.ts:7-23#getRerankWithLog()',
+            '/vscode/src/logged-rerank.ts',
+            7
+        )
+
+        expect(result).toEqual(
+            '[_@vscode/src/logged-rerank.ts:7-23#getRerankWithLog()_](vscode://file/vscode/src/logged-rerank.ts:7) what does this do'
+        )
+    })
+})
diff --git a/lib/shared/src/chat/prompts/display-text.ts b/lib/shared/src/chat/prompts/display-text.ts
@@ -0,0 +1,69 @@
+import { ContextFile } from '../../codebase-context/messages'
+import { ActiveTextEditorSelection } from '../../editor'
+
+/**
+ * Creates display text for the given context files by replacing file names with markdown links.
+ */
+export function createDisplayTextWithFileLinks(files: ContextFile[], text: string): string {
+    let formattedText = text
+    for (const file of files) {
+        if (file?.fileName && file?.uri?.fsPath) {
+            formattedText = replaceFileNameWithMarkdownLink(
+                formattedText,
+                file?.fileName.trim(),
+                file?.uri?.fsPath,
+                file.range?.start?.line
+            )
+        }
+    }
+    return formattedText
+}
+
+/**
+ * Gets the display text to show for the human's input.
+ *
+ * If there is a selection, display the file name + range alongside with human input
+ * If the workspace root is available, it generates a markdown link to the file.
+ */
+export function createDisplayTextWithFileSelection(
+    humanInput: string,
+    selection?: ActiveTextEditorSelection | null
+): string {
+    const fileName = selection?.fileName?.trim()
+    if (!fileName) {
+        return humanInput
+    }
+
+    const displayText = `${humanInput} @${fileName}`
+    const fsPath = selection?.fileUri?.fsPath
+    const startLine = selection?.selectionRange?.start?.line
+    if (!fsPath || !startLine) {
+        return displayText
+    }
+
+    // Create markdown link to the file
+    return replaceFileNameWithMarkdownLink(displayText, `@${fileName}`, fsPath, startLine)
+}
+
+/**
+ * Replaces a file name in given text with markdown link to open that file in editor.
+ * @returns The updated text with the file name replaced by a markdown link.
+ */
+export function replaceFileNameWithMarkdownLink(
+    humanInput: string,
+    fileName: string,
+    fsPath: string,
+    startLine = 0
+): string {
+    // Create markdown link to the file
+    const range = startLine ? `:${startLine}` : ''
+    const fileLink = `vscode://file${fsPath}${range}`
+    const markdownText = `[_${fileName.trim()}_](${fileLink})`
+
+    // Escape special characters in fileName for regex
+    const escapedFileName = fileName.replaceAll(/[$()*+./?[\\\]^{|}-]/g, '\\$&')
+
+    // Updated regex to match the file name with optional line number, range, and symbol name
+    const textToBeReplaced = new RegExp(`(${escapedFileName})(:\\d+(-\\d+)?(#\\S+)?)?(?!\\S)`, 'g')
+    return humanInput.replaceAll(textToBeReplaced, markdownText).trim()
+}
diff --git a/lib/shared/src/chat/prompts/index.ts b/lib/shared/src/chat/prompts/index.ts
@@ -1,3 +1,4 @@
+import { ContextFile } from '../../codebase-context/messages'
 import { ChatEventSource } from '../transcript/messages'
 
 import * as defaultPrompts from './cody.json'
@@ -64,6 +65,9 @@ export interface CodyPrompt {
     type?: CodyPromptType
     slashCommand: string
     mode?: CodyPromptMode
+
+    // internal properties
+    contextFiles?: ContextFile[]
 }
 
 /**
@@ -85,6 +89,7 @@ export interface CodyPromptContext {
     command?: string
     output?: string
     filePath?: string
+    filePaths?: string[]
     directoryPath?: string
     none?: boolean
 }

diff --git a/lib/shared/src/chat/recipes/chat-question.ts b/lib/shared/src/chat/recipes/chat-question.ts
@@ -1,5 +1,10 @@
 import { CodebaseContext } from '../../codebase-context'
-import { ContextMessage, getContextMessageWithResponse } from '../../codebase-context/messages'
+import {
+    ContextFile,
+    ContextMessage,
+    createContextMessageByFile,
+    getContextMessageWithResponse,
+} from '../../codebase-context/messages'
 import { ActiveTextEditorSelection, Editor } from '../../editor'
 import { IntentDetector } from '../../intent-detector'
 import { MAX_CURRENT_FILE_TOKENS, MAX_HUMAN_INPUT_TOKENS } from '../../prompt/constants'
@@ -8,6 +13,7 @@ import {
     populateCurrentEditorSelectedContextTemplate,
 } from '../../prompt/templates'
 import { truncateText } from '../../prompt/truncation'
+import { createDisplayTextWithFileLinks } from '../prompts/display-text'
 import { Interaction } from '../transcript/interaction'
 
 import { isSingleWord, numResults } from './helpers'
@@ -23,17 +29,23 @@ export class ChatQuestion implements Recipe {
         const source = this.id
         const truncatedText = truncateText(humanChatInput, MAX_HUMAN_INPUT_TOKENS)
 
+        const contextFiles = context.userInputContextFiles
+        const displayText = contextFiles?.length
+            ? createDisplayTextWithFileLinks(contextFiles, humanChatInput)
+            : humanChatInput
+
         return Promise.resolve(
             new Interaction(
-                { speaker: 'human', text: truncatedText, displayText: humanChatInput, metadata: { source } },
+                { speaker: 'human', text: truncatedText, displayText, metadata: { source } },
                 { speaker: 'assistant', metadata: { source } },
                 this.getContextMessages(
                     truncatedText,
                     context.editor,
                     context.firstInteraction,
                     context.intentDetector,
                     context.codebaseContext,
-                    context.editor.getActiveTextEditorSelection() || null
+                    context.editor.getActiveTextEditorSelection() || null,
+                    context.userInputContextFiles
                 ),
                 []
             )
@@ -46,7 +58,8 @@ export class ChatQuestion implements Recipe {
         firstInteraction: boolean,
         intentDetector: IntentDetector,
         codebaseContext: CodebaseContext,
-        selection: ActiveTextEditorSelection | null
+        selection: ActiveTextEditorSelection | null,
+        contextFiles?: ContextFile[]
     ): Promise<ContextMessage[]> {
         const contextMessages: ContextMessage[] = []
         // If input is less than 2 words, it means it's most likely a statement or a follow-up question that does not require additional context
@@ -70,6 +83,11 @@ export class ChatQuestion implements Recipe {
             contextMessages.push(...ChatQuestion.getEditorContext(editor))
         }
 
+        if (contextFiles?.length) {
+            const contextFileMessages = await ChatQuestion.getContextFilesContext(editor, contextFiles)
+            contextMessages.push(...contextFileMessages)
+        }
+
         // Add selected text as context when available
         if (selection?.selectedText) {
             contextMessages.push(...ChatQuestion.getEditorSelectionContext(selection))
@@ -97,4 +115,18 @@ export class ChatQuestion implements Recipe {
             selection
         )
     }
+
+    public static async getContextFilesContext(editor: Editor, contextFiles: ContextFile[]): Promise<ContextMessage[]> {
+        const contextFileMessages = []
+        for (const file of contextFiles) {
+            if (file?.uri) {
+                const content = await editor.getTextEditorContentForFile(file?.uri, file.range)
+                if (content) {
+                    const message = createContextMessageByFile(file, content)
+                    contextFileMessages.push(...message)
+                }
+            }
+        }
+        return contextFileMessages
+    }
 }