LLM-enhanced keyword context (#52815)

* Adds `fast` parameter to completions endpoint and `fastChatModel`
param to site config. This is intended for faster chat models that are
useful for simple generations.
* Use the fast chat model to generate a local keyword search. This
replaces the old keyword search mechanism, which stemmed/lemmatized
every word in the user query.
* Use the fast chat model to generate a small set of file fragments to
search for. This is mainly useful for surfacing READMEs for questions
like "What does this project do?"
* Update the set of "files read" presented in the UI to include only
those files actually read into the context window. Previously, we were
showing all files returned by the context fetcher, but in reality, only
a subset of these would fit into the context window.

client/cody-cli/src/app.ts

@@ -111,7 +111,8 @@ async function startCLI() {
         }
     }
 
-    const finalPrompt = await transcript.toPrompt(getPreamble(codebase))
+    const { prompt: finalPrompt, contextFiles } = await transcript.getPromptForLastInteraction(getPreamble(codebase))
+    transcript.setUsedContextFilesForLastInteraction(contextFiles)
 
     let text = ''
     streamCompletions(completionsClient, finalPrompt, {

client/cody-cli/src/context.ts

@@ -1,6 +1,6 @@
 import { CodebaseContext } from '@sourcegraph/cody-shared/src/codebase-context'
 import { SourcegraphEmbeddingsSearchClient } from '@sourcegraph/cody-shared/src/embeddings/client'
-import { KeywordContextFetcher } from '@sourcegraph/cody-shared/src/keyword-context'
+import { KeywordContextFetcher } from '@sourcegraph/cody-shared/src/local-context'
 import { SourcegraphGraphQLAPIClient } from '@sourcegraph/cody-shared/src/sourcegraph-api/graphql'
 import { isError } from '@sourcegraph/cody-shared/src/utils'
 
@@ -26,7 +26,8 @@ export async function createCodebaseContext(
         { useContext: contextType, serverEndpoint },
         codebase,
         embeddingsSearch,
-        new LocalKeywordContextFetcherMock()
+        new LocalKeywordContextFetcherMock(),
+        null
     )
 
     return codebaseContext

client/cody-cli/src/interactions.ts

@@ -45,7 +45,8 @@ export async function interactionFromMessage(
         new Interaction(
             { speaker: 'human', text, displayText: text },
             { speaker: 'assistant', text: '', displayText: '' },
-            contextMessages
+            contextMessages,
+            []
         )
     )
 }

client/cody-shared/package.json

@@ -19,6 +19,10 @@
   },
   "dependencies": {
     "@sourcegraph/common": "workspace:*",
-    "@sourcegraph/http-client": "workspace:*"
+    "@sourcegraph/http-client": "workspace:*",
+    "xml2js": "^0.6.0"
+  },
+  "devDependencies": {
+    "@types/xml2js": "^0.4.11"
   }
 }

client/cody-shared/src/chat/chat.ts

@@ -3,7 +3,9 @@ import { Message } from '../sourcegraph-api'
 import type { SourcegraphCompletionsClient } from '../sourcegraph-api/completions/client'
 import type { CompletionParameters, CompletionCallbacks } from '../sourcegraph-api/completions/types'
 
-const DEFAULT_CHAT_COMPLETION_PARAMETERS: Omit<CompletionParameters, 'messages'> = {
+type ChatParameters = Omit<CompletionParameters, 'messages'>
+
+const DEFAULT_CHAT_COMPLETION_PARAMETERS: ChatParameters = {
     temperature: 0.2,
     maxTokensToSample: SOLUTION_TOKEN_LENGTH,
     topK: -1,
@@ -13,10 +15,17 @@ const DEFAULT_CHAT_COMPLETION_PARAMETERS: Omit<CompletionParameters, 'messages'>
 export class ChatClient {
     constructor(private completions: SourcegraphCompletionsClient) {}
 
-    public chat(messages: Message[], cb: CompletionCallbacks): () => void {
+    public chat(messages: Message[], cb: CompletionCallbacks, params?: Partial<ChatParameters>): () => void {
         const isLastMessageFromHuman = messages.length > 0 && messages[messages.length - 1].speaker === 'human'
         const augmentedMessages = isLastMessageFromHuman ? messages.concat([{ speaker: 'assistant' }]) : messages
 
-        return this.completions.stream({ messages: augmentedMessages, ...DEFAULT_CHAT_COMPLETION_PARAMETERS }, cb)
+        return this.completions.stream(
+            {
+                ...DEFAULT_CHAT_COMPLETION_PARAMETERS,
+                ...params,
+                messages: augmentedMessages,
+            },
+            cb
+        )
     }
 }

client/cody-shared/src/chat/client.ts

@@ -69,7 +69,7 @@ export async function createClient({
 
     const embeddingsSearch = repoId ? new SourcegraphEmbeddingsSearchClient(graphqlClient, repoId, true) : null
 
-    const codebaseContext = new CodebaseContext(config, config.codebase, embeddingsSearch, null)
+    const codebaseContext = new CodebaseContext(config, config.codebase, embeddingsSearch, null, null)
 
     const intentDetector = new SourcegraphIntentDetectorClient(graphqlClient)
 
@@ -116,10 +116,11 @@ export async function createClient({
 
         sendTranscript()
 
-        const prompt = await transcript.toPrompt(getPreamble(config.codebase))
+        const { prompt, contextFiles } = await transcript.getPromptForLastInteraction(getPreamble(config.codebase))
+        transcript.setUsedContextFilesForLastInteraction(contextFiles)
+
         const responsePrefix = interaction.getAssistantMessage().prefix ?? ''
         let rawText = ''
-
         chatClient.chat(prompt, {
             onChange(_rawText) {
                 rawText = _rawText

client/cody-shared/src/chat/recipes/chat-question.ts

@@ -31,7 +31,8 @@ export class ChatQuestion implements Recipe {
                     context.intentDetector,
                     context.codebaseContext,
                     context.editor.getActiveTextEditorSelection() || null
-                )
+                ),
+                []
             )
         )
     }

client/cody-shared/src/chat/recipes/context-search.ts

@@ -46,7 +46,8 @@ export class ContextSearch implements Recipe {
                 text: '',
                 displayText: await this.displaySearchResults(truncatedText, context.codebaseContext, wsRootPath),
             },
-            new Promise(resolve => resolve([]))
+            new Promise(resolve => resolve([])),
+            []
         )
     }
 

client/cody-shared/src/chat/recipes/explain-code-detailed.ts

@@ -32,7 +32,8 @@ export class ExplainCodeDetailed implements Recipe {
                 truncatedFollowingText,
                 selection,
                 context.codebaseContext
-            )
+            ),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/explain-code-high-level.ts

@@ -32,7 +32,8 @@ export class ExplainCodeHighLevel implements Recipe {
                 truncatedFollowingText,
                 selection,
                 context.codebaseContext
-            )
+            ),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/file-touch.ts

@@ -103,7 +103,8 @@ export class FileTouch implements Recipe {
                     speaker: 'assistant',
                     prefix: 'Working on it! I will notify you when the file is ready.\n',
                 },
-                this.getContextMessages(selection, currentDir)
+                this.getContextMessages(selection, currentDir),
+                []
             )
         )
     }

client/cody-shared/src/chat/recipes/find-code-smells.ts

@@ -36,7 +36,8 @@ If you have no ideas because the code looks fine, feel free to say that it alrea
                 prefix: assistantResponsePrefix,
                 text: assistantResponsePrefix,
             },
-            new Promise(resolve => resolve([]))
+            new Promise(resolve => resolve([])),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/fixup.ts

@@ -67,7 +67,8 @@ export class Fixup implements Recipe {
                     speaker: 'assistant',
                     prefix: 'Check your document for updates from Cody.\n',
                 },
-                this.getContextMessages(selection.selectedText, context.codebaseContext)
+                this.getContextMessages(selection.selectedText, context.codebaseContext),
+                []
             )
         )
     }

client/cody-shared/src/chat/recipes/generate-docstring.ts

@@ -60,7 +60,8 @@ export class GenerateDocstring implements Recipe {
                 truncatedFollowingText,
                 selection,
                 context.codebaseContext
-            )
+            ),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/generate-pr-description.ts

@@ -52,7 +52,8 @@ export class PrDescription implements Recipe {
                     prefix: emptyGitCommitMessage,
                     text: emptyGitCommitMessage,
                 },
-                Promise.resolve([])
+                Promise.resolve([]),
+                []
             )
         }
 
@@ -71,7 +72,8 @@ export class PrDescription implements Recipe {
                 prefix: assistantResponsePrefix,
                 text: assistantResponsePrefix,
             },
-            Promise.resolve([])
+            Promise.resolve([]),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/generate-release-notes.ts

@@ -71,7 +71,8 @@ export class ReleaseNotes implements Recipe {
                     prefix: emptyGitLogMessage,
                     text: emptyGitLogMessage,
                 },
-                Promise.resolve([])
+                Promise.resolve([]),
+                []
             )
         }
 
@@ -91,7 +92,8 @@ export class ReleaseNotes implements Recipe {
                 prefix: assistantResponsePrefix,
                 text: assistantResponsePrefix,
             },
-            Promise.resolve([])
+            Promise.resolve([]),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/generate-test.ts

@@ -44,7 +44,8 @@ export class GenerateTest implements Recipe {
                 truncatedFollowingText,
                 selection,
                 context.codebaseContext
-            )
+            ),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/git-log.ts

@@ -65,7 +65,8 @@ export class GitHistory implements Recipe {
                     prefix: emptyGitLogMessage,
                     text: emptyGitLogMessage,
                 },
-                Promise.resolve([])
+                Promise.resolve([]),
+                []
             )
         }
 
@@ -84,7 +85,8 @@ export class GitHistory implements Recipe {
                 prefix: assistantResponsePrefix,
                 text: assistantResponsePrefix,
             },
-            Promise.resolve([])
+            Promise.resolve([]),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/improve-variable-names.ts

@@ -44,7 +44,8 @@ export class ImproveVariableNames implements Recipe {
                 truncatedFollowingText,
                 selection,
                 context.codebaseContext
-            )
+            ),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/inline-chat.ts

@@ -56,7 +56,8 @@ export class InlineAssist implements Recipe {
                     displayText,
                 },
                 { speaker: 'assistant' },
-                this.getContextMessages(truncatedText, context.codebaseContext, selection, context.editor)
+                this.getContextMessages(truncatedText, context.codebaseContext, selection, context.editor),
+                []
             )
         )
     }

client/cody-shared/src/chat/recipes/next-questions.ts

@@ -31,7 +31,8 @@ export class NextQuestions implements Recipe {
                     prefix: assistantResponsePrefix,
                     text: assistantResponsePrefix,
                 },
-                this.getContextMessages(promptMessage, context.editor, context.intentDetector, context.codebaseContext)
+                this.getContextMessages(promptMessage, context.editor, context.intentDetector, context.codebaseContext),
+                []
             )
         )
     }

client/cody-shared/src/chat/recipes/non-stop.ts

@@ -76,7 +76,8 @@ export class NonStop implements Recipe {
                     speaker: 'assistant',
                     prefix: 'Check your document for updates from Cody.',
                 },
-                this.getContextMessages(selection.selectedText, context.codebaseContext)
+                this.getContextMessages(selection.selectedText, context.codebaseContext),
+                []
             )
         )
     }

client/cody-shared/src/chat/recipes/optimize-code.ts

@@ -51,7 +51,8 @@ However if no optimization is possible; just say the code is already optimized.
                 truncatedFollowingText,
                 selection,
                 context.codebaseContext
-            )
+            ),
+            []
         )
     }
 }

client/cody-shared/src/chat/recipes/translate.ts

@@ -39,7 +39,8 @@ export class TranslateToLanguage implements Recipe {
                 prefix: assistantResponsePrefix,
                 text: assistantResponsePrefix,
             },
-            Promise.resolve([])
+            Promise.resolve([]),
+            []
         )
     }
 }