cody: update local context fetching

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-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 ? 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/transcript/index.ts

@@ -1,5 +1,6 @@
-import { OldContextMessage } from '../../codebase-context/messages'
+import { ContextFile, ContextMessage, OldContextMessage } from '../../codebase-context/messages'
 import { CHARS_PER_TOKEN, MAX_AVAILABLE_PROMPT_LENGTH } from '../../prompt/constants'
+import { PromptMixin } from '../../prompt/prompt-mixin'
 import { Message } from '../../sourcegraph-api'
 
 import { Interaction, InteractionJSON } from './interaction'
@@ -20,18 +21,19 @@ export interface TranscriptJSON {
 }
 
 /**
- * A transcript of a conversation between a human and an assistant.
+ * The "model" class that tracks the call and response of the Cody chat box.
+ * Any "controller" logic belongs outside of this class.
  */
 export class Transcript {
     public static fromJSON(json: TranscriptJSON): Transcript {
         return new Transcript(
             json.interactions.map(
-                ({ humanMessage, assistantMessage, context, timestamp }) =>
+                ({ humanMessage, assistantMessage, fullContext, usedContextFiles, timestamp }) =>
                     new Interaction(
                         humanMessage,
                         assistantMessage,
                         Promise.resolve(
-                            context.map(message => {
+                            fullContext.map(message => {
                                 if (message.file) {
                                     return message
                                 }
@@ -44,6 +46,7 @@ export class Transcript {
                                 return message
                             })
                         ),
+                        usedContextFiles,
                         timestamp || new Date().toISOString()
                     )
             ),
@@ -144,20 +147,50 @@ export class Transcript {
         return -1
     }
 
-    public async toPrompt(preamble: Message[] = []): Promise<Message[]> {
+    public async getPromptForLastInteraction(
+        preamble: Message[] = []
+    ): Promise<{ prompt: Message[]; contextFiles: ContextFile[] }> {
+        if (this.interactions.length == 0) {
+            return { prompt: [], contextFiles: [] }
+        }
+
         const lastInteractionWithContextIndex = await this.getLastInteractionWithContextIndex()
         const messages: Message[] = []
         for (let index = 0; index < this.interactions.length; index++) {
-            // Include context messages for the last interaction that has a non-empty context.
-            const interactionMessages = await this.interactions[index].toPrompt(
-                index === lastInteractionWithContextIndex
-            )
-            messages.push(...interactionMessages)
+            const interaction = this.interactions[index]
+            const humanMessage = PromptMixin.mixInto(interaction.getHumanMessage())
+            const assistantMessage = interaction.getAssistantMessage()
+            const contextMessages = await interaction.getFullContext()
+            if (index === lastInteractionWithContextIndex) {
+                messages.push(...contextMessages, humanMessage, assistantMessage)
+            } else {
+                messages.push(humanMessage, assistantMessage)
+            }
         }
 
         const preambleTokensUsage = preamble.reduce((acc, message) => acc + estimateTokensUsage(message), 0)
         const truncatedMessages = truncatePrompt(messages, MAX_AVAILABLE_PROMPT_LENGTH - preambleTokensUsage)
-        return [...preamble, ...truncatedMessages]
+
+        // Return what context fits in the window
+        const contextFiles: ContextFile[] = []
+        for (const msg of truncatedMessages) {
+            const contextFile = (msg as ContextMessage).file
+            if (contextFile) {
+                contextFiles.push(contextFile)
+            }
+        }
+
+        return {
+            prompt: [...preamble, ...truncatedMessages],
+            contextFiles,
+        }
+    }
+
+    public async setUsedContextFilesForLastInteraction(contextFiles: ContextFile[]) {
+        if (this.interactions.length === 0) {
+            throw new Error('Cannot set context files for empty transcript')
+        }
+        this.interactions[this.interactions.length - 1].setUsedContext(contextFiles)
     }
 
     public toChat(): ChatMessage[] {

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

@@ -1,97 +1,68 @@
 import { ContextMessage, ContextFile } from '../../codebase-context/messages'
-import { PromptMixin } from '../../prompt/prompt-mixin'
-import { Message } from '../../sourcegraph-api'
 
 import { ChatMessage, InteractionMessage } from './messages'
 
 export interface InteractionJSON {
     humanMessage: InteractionMessage
     assistantMessage: InteractionMessage
-    context: ContextMessage[]
+    fullContext: ContextMessage[]
+    usedContextFiles: ContextFile[]
     timestamp: string
 }
 
 export class Interaction {
-    private readonly humanMessage: InteractionMessage
-    private assistantMessage: InteractionMessage
-    private cachedContextFiles: ContextFile[] = []
-    public readonly timestamp: string
-    private readonly context: Promise<ContextMessage[]>
-
     constructor(
-        humanMessage: InteractionMessage,
-        assistantMessage: InteractionMessage,
-        context: Promise<ContextMessage[]>,
-        timestamp: string = new Date().toISOString()
-    ) {
-        this.humanMessage = humanMessage
-        this.assistantMessage = assistantMessage
-        this.timestamp = timestamp
-
-        // This is some hacky behavior: returns a promise that resolves to the same array that was passed,
-        // but also caches the context file names in memory as a side effect.
-        this.context = context.then(messages => {
-            const contextFilesMap = messages.reduce((map, { file }) => {
-                if (!file?.fileName) {
-                    return map
-                }
-                map[`${file.repoName || 'repo'}@${file?.revision || 'HEAD'}/${file.fileName}`] = file
-                return map
-            }, {} as { [key: string]: ContextFile })
-
-            // Cache the context files so we don't have to block the UI when calling `toChat` by waiting for the context to resolve.
-            this.cachedContextFiles = [
-                ...Object.keys(contextFilesMap)
-                    .sort((a, b) => a.localeCompare(b))
-                    .map((key: string) => contextFilesMap[key]),
-            ]
-
-            return messages
-        })
-    }
+        private readonly humanMessage: InteractionMessage,
+        private assistantMessage: InteractionMessage,
+        private fullContext: Promise<ContextMessage[]>,
+        private usedContextFiles: ContextFile[],
+        public readonly timestamp: string = new Date().toISOString()
+    ) {}
 
     public getAssistantMessage(): InteractionMessage {
-        return this.assistantMessage
+        return { ...this.assistantMessage }
     }
 
     public setAssistantMessage(assistantMessage: InteractionMessage): void {
         this.assistantMessage = assistantMessage
     }
 
+    public getHumanMessage(): InteractionMessage {
+        return { ...this.humanMessage }
+    }
+
+    public async getFullContext(): Promise<ContextMessage[]> {
+        const msgs = await this.fullContext
+        return msgs.map(msg => ({ ...msg }))
+    }
+
     public async hasContext(): Promise<boolean> {
-        const contextMessages = await this.context
+        const contextMessages = await this.fullContext
         return contextMessages.length > 0
     }
 
-    public async toPrompt(includeContext: boolean): Promise<Message[]> {
-        const messages: (ContextMessage | InteractionMessage)[] = [
-            PromptMixin.mixInto(this.humanMessage),
-            this.assistantMessage,
-        ]
-        if (includeContext) {
-            messages.unshift(...(await this.context))
-        }
-
-        return messages.map(message => ({ speaker: message.speaker, text: message.text }))
+    public setUsedContext(usedContextFiles: ContextFile[]): void {
+        this.usedContextFiles = usedContextFiles
     }
 
     /**
      * Converts the interaction to chat message pair: one message from a human, one from an assistant.
      */
     public toChat(): ChatMessage[] {
-        return [this.humanMessage, { ...this.assistantMessage, contextFiles: this.cachedContextFiles }]
+        return [this.humanMessage, { ...this.assistantMessage, contextFiles: this.usedContextFiles }]
     }
 
     public async toChatPromise(): Promise<ChatMessage[]> {
-        await this.context
+        await this.fullContext
         return this.toChat()
     }
 
     public async toJSON(): Promise<InteractionJSON> {
         return {
             humanMessage: this.humanMessage,
             assistantMessage: this.assistantMessage,
-            context: await this.context,
+            fullContext: await this.fullContext,
+            usedContextFiles: this.usedContextFiles,
             timestamp: this.timestamp,
         }
     }