diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index a8b602944dd..16c37cabe97 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts @@ -29,7 +29,7 @@ import { ViewDiff, AcceptDiff, QuickCommandGroupActionClick, - MergedRelevantDocument, + DocumentReference, FileClick, RelevantTextDocumentAddition, } from './model' @@ -974,53 +974,53 @@ export class ChatController { } const relativePaths = await this.resolveContextCommandPayload(triggerPayload) - // TODO: resolve the context into real context up to 90k triggerPayload.useRelevantDocuments = false - triggerPayload.mergedRelevantDocuments = [] - if (triggerPayload.message) { - triggerPayload.useRelevantDocuments = triggerPayload.context?.some( - (context) => typeof context !== 'string' && context.command === '@workspace' - ) - if (triggerPayload.useRelevantDocuments) { - triggerPayload.message = triggerPayload.message.replace(/workspace/, '') - if (CodeWhispererSettings.instance.isLocalIndexEnabled()) { - const start = performance.now() - triggerPayload.relevantTextDocuments = await LspController.instance.query(triggerPayload.message) - triggerPayload.mergedRelevantDocuments = this.mergeRelevantTextDocuments( - triggerPayload.relevantTextDocuments + triggerPayload.documentReferences = [] + triggerPayload.useRelevantDocuments = triggerPayload.context?.some( + (context) => typeof context !== 'string' && context.command === '@workspace' + ) + if (triggerPayload.useRelevantDocuments && triggerPayload.message) { + triggerPayload.message = triggerPayload.message.replace(/workspace/, '') + if (CodeWhispererSettings.instance.isLocalIndexEnabled()) { + const start = performance.now() + triggerPayload.relevantTextDocuments = await LspController.instance.query(triggerPayload.message) + + for (const doc of triggerPayload.relevantTextDocuments) { + getLogger().info( + `amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}, start line: ${doc.startLine}, end line: ${doc.endLine}` ) - for (const doc of triggerPayload.relevantTextDocuments) { - getLogger().info( - `amazonq: Using workspace files ${doc.relativeFilePath}, content(partial): ${doc.text?.substring(0, 200)}, start line: ${doc.startLine}, end line: ${doc.endLine}` - ) - } - triggerPayload.projectContextQueryLatencyMs = performance.now() - start - } else { - this.messenger.sendOpenSettingsMessage(triggerID, tabID) - return } + triggerPayload.projectContextQueryLatencyMs = performance.now() - start + } else { + this.messenger.sendOpenSettingsMessage(triggerID, tabID) + return } } + triggerPayload.documentReferences = this.mergeRelevantTextDocuments(triggerPayload.relevantTextDocuments || []) + + // TODO: make sure the user input + current focused document + relevantDocument + additionalContext + // combined does not exceed 100k characters before generating the request payload. + // Do truncation and make sure triggerPayload.documentReferences is up-to-date after truncation const request = triggerPayloadToChatRequest(triggerPayload) const session = this.sessionStorage.getSession(tabID) session.currentContextId++ session.contexts.set(session.currentContextId, new Map()) - if (triggerPayload.mergedRelevantDocuments !== undefined) { - const relativePathsOfMergedRelevantDocuments = triggerPayload.mergedRelevantDocuments.map( + if (triggerPayload.documentReferences !== undefined) { + const relativePathsOfMergedRelevantDocuments = triggerPayload.documentReferences.map( (doc) => doc.relativeFilePath ) for (const relativePath of relativePaths) { if (!relativePathsOfMergedRelevantDocuments.includes(relativePath)) { - triggerPayload.mergedRelevantDocuments.push({ + triggerPayload.documentReferences.push({ relativeFilePath: relativePath, lineRanges: [{ first: -1, second: -1 }], }) } } - if (triggerPayload.mergedRelevantDocuments) { - for (const doc of triggerPayload.mergedRelevantDocuments) { + if (triggerPayload.documentReferences) { + for (const doc of triggerPayload.documentReferences) { const currentContext = session.contexts.get(session.currentContextId) if (currentContext) { currentContext.set(doc.relativeFilePath, doc.lineRanges) @@ -1037,7 +1037,7 @@ export class ChatController { let response: MessengerResponseType | undefined = undefined session.createNewTokenSource() try { - this.messenger.sendInitalStream(tabID, triggerID, triggerPayload.mergedRelevantDocuments) + this.messenger.sendInitalStream(tabID, triggerID, triggerPayload.documentReferences) this.telemetryHelper.setConversationStreamStartTime(tabID) if (isSsoConnection(AuthUtil.instance.conn)) { const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request) @@ -1068,11 +1068,9 @@ export class ChatController { } } - private mergeRelevantTextDocuments( - documents: RelevantTextDocumentAddition[] | undefined - ): MergedRelevantDocument[] | undefined { - if (documents === undefined) { - return undefined + private mergeRelevantTextDocuments(documents: RelevantTextDocumentAddition[]): DocumentReference[] { + if (documents.length === 0) { + return [] } return Object.entries( documents.reduce>((acc, doc) => { diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts index dc4a5074fc4..0f434c6dc6d 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts @@ -25,7 +25,7 @@ import { ChatMessage, ErrorMessage, FollowUp, Suggestion } from '../../../view/c import { ChatSession } from '../../../clients/chat/v0/chat' import { ChatException } from './model' import { CWCTelemetryHelper } from '../telemetryHelper' -import { ChatPromptCommandType, MergedRelevantDocument, TriggerPayload } from '../model' +import { ChatPromptCommandType, DocumentReference, TriggerPayload } from '../model' import { getHttpStatusCode, getRequestId, ToolkitError } from '../../../../shared/errors' import { keys } from '../../../../shared/utilities/tsUtils' import { getLogger } from '../../../../shared/logger/logger' @@ -69,7 +69,7 @@ export class Messenger { public sendInitalStream( tabID: string, triggerID: string, - mergedRelevantDocuments: MergedRelevantDocument[] | undefined + mergedRelevantDocuments: DocumentReference[] | undefined ) { this.dispatcher.sendChatMessage( new ChatMessage( diff --git a/packages/core/src/codewhispererChat/controllers/chat/model.ts b/packages/core/src/codewhispererChat/controllers/chat/model.ts index 72b27a1c43b..f8033649599 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/model.ts +++ b/packages/core/src/codewhispererChat/controllers/chat/model.ts @@ -189,7 +189,9 @@ export interface TriggerPayload { readonly context?: string[] | QuickActionCommand[] relevantTextDocuments?: RelevantTextDocumentAddition[] additionalContents?: AdditionalContentEntry[] - mergedRelevantDocuments?: MergedRelevantDocument[] + // a reference to all documents used in chat payload + // for providing better context transparency + documentReferences?: DocumentReference[] useRelevantDocuments?: boolean traceId?: string } @@ -197,7 +199,7 @@ export interface TriggerPayload { // TODO move this to API definition (or just use this across the codebase) export type RelevantTextDocumentAddition = RelevantTextDocument & { startLine: number; endLine: number } -export interface MergedRelevantDocument { +export interface DocumentReference { readonly relativeFilePath: string readonly lineRanges: Array<{ first: number; second: number }> } diff --git a/packages/core/src/codewhispererChat/view/connector/connector.ts b/packages/core/src/codewhispererChat/view/connector/connector.ts index 175cfbd17ad..0b2b29498c4 100644 --- a/packages/core/src/codewhispererChat/view/connector/connector.ts +++ b/packages/core/src/codewhispererChat/view/connector/connector.ts @@ -8,7 +8,7 @@ import { MessagePublisher } from '../../../amazonq/messages/messagePublisher' import { EditorContextCommandType } from '../../commands/registerCommands' import { AuthFollowUpType } from '../../../amazonq/auth/model' import { ChatItemButton, ChatItemFormItem, MynahUIDataModel, QuickActionCommand } from '@aws/mynah-ui' -import { MergedRelevantDocument } from '../../controllers/chat/model' +import { DocumentReference } from '../../controllers/chat/model' class UiMessage { readonly time: number = Date.now() @@ -207,7 +207,7 @@ export interface ChatMessageProps { readonly messageID: string readonly userIntent: string | undefined readonly codeBlockLanguage: string | undefined - readonly contextList: MergedRelevantDocument[] | undefined + readonly contextList: DocumentReference[] | undefined } export class ChatMessage extends UiMessage { @@ -222,7 +222,7 @@ export class ChatMessage extends UiMessage { readonly messageID: string | undefined readonly userIntent: string | undefined readonly codeBlockLanguage: string | undefined - readonly contextList: MergedRelevantDocument[] | undefined + readonly contextList: DocumentReference[] | undefined override type = 'chatMessage' constructor(props: ChatMessageProps, tabID: string) {