-
Notifications
You must be signed in to change notification settings - Fork 1.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
LLM-enhanced keyword context #52815
LLM-enhanced keyword context #52815
Changes from 3 commits
8959f62
b76fc86
3d69d7f
d80de1f
e14f183
b16a27b
0b12296
16d1eef
e1f8075
764829d
1724f50
15952a8
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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. | ||
*/ | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We simplify the |
||
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[] { | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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' | ||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Likewise, we also make the |
||
export interface InteractionJSON { | ||
humanMessage: InteractionMessage | ||
assistantMessage: InteractionMessage | ||
context: ContextMessage[] | ||
fullContext: ContextMessage[] | ||
usedContextFiles: ContextFile[] | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Instead of having
|
||
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, | ||
} | ||
} | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pass through the underlying completions parameters, so we can set things like temperature.