Skip to content

Commit

Permalink
refactor(cw): Rename the mergedTextDocument field. (#6666)
Browse files Browse the repository at this point in the history
## Problem
The mergedRelevantTextDocument should have a better name. It is actually
the documents that are referenced.

## Solution


---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.
  • Loading branch information
leigaol authored Feb 24, 2025
1 parent 3130c68 commit ae20598
Show file tree
Hide file tree
Showing 4 changed files with 41 additions and 41 deletions.
66 changes: 32 additions & 34 deletions packages/core/src/codewhispererChat/controllers/chat/controller.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ import {
ViewDiff,
AcceptDiff,
QuickCommandGroupActionClick,
MergedRelevantDocument,
DocumentReference,
FileClick,
RelevantTextDocumentAddition,
} from './model'
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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<Record<string, { first: number; second: number }[]>>((acc, doc) => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -69,7 +69,7 @@ export class Messenger {
public sendInitalStream(
tabID: string,
triggerID: string,
mergedRelevantDocuments: MergedRelevantDocument[] | undefined
mergedRelevantDocuments: DocumentReference[] | undefined
) {
this.dispatcher.sendChatMessage(
new ChatMessage(
Expand Down
6 changes: 4 additions & 2 deletions packages/core/src/codewhispererChat/controllers/chat/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -189,15 +189,17 @@ 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
}

// 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 }>
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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 {
Expand All @@ -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) {
Expand Down

0 comments on commit ae20598

Please sign in to comment.