custom command: update get context for current file

lib/shared/src/chat/prompts/templates.ts

@@ -12,6 +12,7 @@ const prevent_hallucinations =
 export const answers = {
     terminal: 'Noted. I will answer your next question based on this terminal output with the code you just shared.',
     selection: 'Noted. I will refer to this code you selected in the editor to answer your question.',
+    file: 'Noted. I will refer to this file you are looking at, with your selected code inside the <selected> tags to answer your question.',
 }
 
 export const prompts = {

lib/shared/src/chat/prompts/vscode-context.ts

@@ -9,6 +9,7 @@ import { MAX_CURRENT_FILE_TOKENS } from '../../prompt/constants'
 import {
     populateCodeContextTemplate,
     populateCurrentEditorContextTemplate,
+    populateCurrentFileFromEditorSelectionContextTemplate,
     populateCurrentSelectedCodeContextTemplate,
     populateTerminalOutputContextTemplate,
 } from '../../prompt/templates'
@@ -346,3 +347,18 @@ export function getHumanDisplayTextWithFileName(
     const fileLink = `vscode://file${fileUri.fsPath}:${startLineNumber}`
     return `${humanInput}\n\nFile: [_${fileName}:${fileRange}_](${fileLink})`
 }
+
+/**
+ * Gets context messages for the current open file's content based on the editor selection if any (or use visible content)
+ */
+export function getCurrentFileContextFromEditorSelection(selection: ActiveTextEditorSelection): ContextMessage[] {
+    if (!selection.selectedText) {
+        return []
+    }
+
+    return getContextMessageWithResponse(
+        populateCurrentFileFromEditorSelectionContextTemplate(selection, selection.fileName),
+        selection,
+        answers.file
+    )
+}

lib/shared/src/chat/recipes/custom-prompt.ts

@@ -1,14 +1,8 @@
 import { CodebaseContext } from '../../codebase-context'
 import { ContextMessage } from '../../codebase-context/messages'
 import { ActiveTextEditorSelection, Editor } from '../../editor'
-import {
-    MAX_HUMAN_INPUT_TOKENS,
-    MAX_RECIPE_INPUT_TOKENS,
-    MAX_RECIPE_SURROUNDING_TOKENS,
-    NUM_CODE_RESULTS,
-    NUM_TEXT_RESULTS,
-} from '../../prompt/constants'
-import { truncateText, truncateTextStart } from '../../prompt/truncation'
+import { MAX_HUMAN_INPUT_TOKENS, NUM_CODE_RESULTS, NUM_TEXT_RESULTS } from '../../prompt/constants'
+import { truncateText } from '../../prompt/truncation'
 import { CodyPromptContext, defaultCodyPromptContext } from '../prompts'
 import { prompts } from '../prompts/templates'
 import {
@@ -19,6 +13,7 @@ import {
 } from '../prompts/utils'
 import {
     getCurrentDirContext,
+    getCurrentFileContextFromEditorSelection,
     getEditorDirContext,
     getEditorOpenTabsContext,
     getEditorSelectionContext,
@@ -29,7 +24,7 @@ import {
 } from '../prompts/vscode-context'
 import { Interaction } from '../transcript/interaction'
 
-import { getContextMessagesFromSelection, getFileExtension, getNormalizedLanguageName, numResults } from './helpers'
+import { getFileExtension, getNormalizedLanguageName, numResults } from './helpers'
 import { Recipe, RecipeContext, RecipeID } from './recipe'
 
 /** ======================================================
@@ -160,16 +155,7 @@ export class CustomPrompt implements Recipe {
         // If currentFile is true, or when selection is true but there is no selected text
         // then we want to include the current file context
         if (selection && (isContextRequired.currentFile || (isContextRequired.selection && !selection?.selectedText))) {
-            const truncatedSelectedText = truncateText(selection.selectedText, MAX_RECIPE_INPUT_TOKENS)
-            const truncatedPrecedingText = truncateTextStart(selection.precedingText, MAX_RECIPE_SURROUNDING_TOKENS)
-            const truncatedFollowingText = truncateText(selection.followingText, MAX_RECIPE_SURROUNDING_TOKENS)
-            const contextMsg = await getContextMessagesFromSelection(
-                truncatedSelectedText,
-                truncatedPrecedingText,
-                truncatedFollowingText,
-                selection,
-                codebaseContext
-            )
+            const contextMsg = getCurrentFileContextFromEditorSelection(selection)
             currentFileContextStack.push(...contextMsg)
         }
 

lib/shared/src/prompt/templates.ts

@@ -1,7 +1,10 @@
 import path from 'path'
 
 import { getFileExtension, getNormalizedLanguageName } from '../chat/recipes/helpers'
-import { ActiveTextEditorDiagnostic } from '../editor'
+import { ActiveTextEditorDiagnostic, ActiveTextEditorSelection } from '../editor'
+
+import { MAX_CURRENT_FILE_TOKENS, MAX_RECIPE_INPUT_TOKENS, MAX_RECIPE_SURROUNDING_TOKENS } from './constants'
+import { truncateText, truncateTextStart } from './truncation'
 
 const CODE_CONTEXT_TEMPLATE = `Use following code snippet from file \`{filePath}\`:
 \`\`\`{language}
@@ -132,3 +135,27 @@ export function populateCurrentSelectedCodeContextTemplate(code: string, filePat
         .replace(/{filePath}/g, filePath)
         .replace('{languageName}', languageName)
 }
+
+const CURRENT_FILE_CONTEXT_TEMPLATE = `"This is the {languageName} file \`{filePath}\` I am looking at, with my selected code in <selected> tags:
+{precedingText}<selected>
+{selectedText}
+</selected>{followingText}`
+
+export function populateCurrentFileFromEditorSelectionContextTemplate(
+    selection: ActiveTextEditorSelection,
+    filePath: string
+): string {
+    const extension = getFileExtension(filePath)
+    const languageName = getNormalizedLanguageName(extension)
+    const truncatedSelectedText = truncateText(selection.selectedText, MAX_RECIPE_INPUT_TOKENS) || ''
+    const truncatedPrecedingText = truncateTextStart(selection.precedingText, MAX_RECIPE_SURROUNDING_TOKENS)
+    const truncatedFollowingText = truncateText(selection.followingText, MAX_RECIPE_SURROUNDING_TOKENS)
+
+    const fileContext = CURRENT_FILE_CONTEXT_TEMPLATE.replace('{languageName}', languageName)
+        .replace(/{filePath}/g, filePath)
+        .replace('{followingText}', truncatedFollowingText)
+        .replace('{selectedText}', truncatedSelectedText)
+        .replace('{precedingText}', truncatedPrecedingText)
+
+    return truncateText(fileContext, MAX_CURRENT_FILE_TOKENS)
+}

vscode/CHANGELOG.md

@@ -14,6 +14,7 @@ Starting from `0.2.0`, Cody is using `major.EVEN_NUMBER.patch` for release versi
 ### Fixed
 
 - Inline Chat: Fix issue where state was not being set correctly, causing Cody Commands to use the selection range from the last created Inline Chat instead of the current selection. [pull/602](https://github.com/sourcegraph/cody/pull/602)
+- Cody Commands: Commands that use the current file as context now correctly generate context message for the current file instead of using codebase context generated from current selection. [pull/683](https://github.com/sourcegraph/cody/pull/683)
 
 ### Changed