Skip to content

Commit 788920b

Browse files
yzhangokyzhangok
authored
fix(amazonq): add image context to chat history (aws#1859)
* fix(amazonq): add image context to chat history * fix(amazonq): remove un-used import * fix(amazonq): truncate image blocks * fix(amazonq): updaterequestinput reset * fix(amazonq): add image context to additionalContext * fix(amazonq): rebase * fix(amazonq): add comments * fix(amazonq): fix test * fix(amazonq): comments * fix(amazonq): change variable name * fix(amazonq): change variable name * fix(amazonq): change variable name
1 parent 3c0592f commit 788920b

File tree

7 files changed

+349
-60
lines changed

7 files changed

+349
-60
lines changed

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -491,6 +491,7 @@ describe('AgenticChatController', () => {
491491
{
492492
userInputMessage: {
493493
content: 'Previous question',
494+
images: [],
494495
origin: 'IDE',
495496
userInputMessageContext: { toolResults: [] },
496497
userIntent: undefined,
@@ -540,6 +541,7 @@ describe('AgenticChatController', () => {
540541
{
541542
userInputMessage: {
542543
content: 'Previous question',
544+
images: [],
543545
origin: 'IDE',
544546
userInputMessageContext: { toolResults: [] },
545547
userIntent: undefined,
@@ -1803,6 +1805,117 @@ describe('AgenticChatController', () => {
18031805
assert.strictEqual(request.conversationState?.history?.length || 0, 3)
18041806
assert.strictEqual(result, 298000)
18051807
})
1808+
1809+
it('should truncate images when they exceed budget', () => {
1810+
const request: GenerateAssistantResponseCommandInput = {
1811+
conversationState: {
1812+
currentMessage: {
1813+
userInputMessage: {
1814+
content: 'a'.repeat(493_400),
1815+
images: [
1816+
{
1817+
format: 'png',
1818+
source: {
1819+
bytes: new Uint8Array(1000), // 3.3 chars
1820+
},
1821+
},
1822+
{
1823+
format: 'png',
1824+
source: {
1825+
bytes: new Uint8Array(2000000), //6600 chars - should be removed
1826+
},
1827+
},
1828+
{
1829+
format: 'png',
1830+
source: {
1831+
bytes: new Uint8Array(1000), // 3.3 chars
1832+
},
1833+
},
1834+
],
1835+
},
1836+
},
1837+
chatTriggerType: undefined,
1838+
},
1839+
}
1840+
const result = chatController.truncateRequest(request)
1841+
1842+
// Should only keep the first and third images (small ones)
1843+
assert.strictEqual(request.conversationState?.currentMessage?.userInputMessage?.images?.length, 2)
1844+
assert.strictEqual(result, 500000 - 493400 - 3.3 - 3.3) // remaining budget after content and images
1845+
})
1846+
1847+
it('should handle images without bytes', () => {
1848+
const request: GenerateAssistantResponseCommandInput = {
1849+
conversationState: {
1850+
currentMessage: {
1851+
userInputMessage: {
1852+
content: 'a'.repeat(400_000),
1853+
images: [
1854+
{
1855+
format: 'png',
1856+
source: {
1857+
bytes: null as any,
1858+
},
1859+
},
1860+
{
1861+
format: 'png',
1862+
source: {
1863+
bytes: new Uint8Array(1000), // 3.3 chars
1864+
},
1865+
},
1866+
],
1867+
},
1868+
},
1869+
chatTriggerType: undefined,
1870+
},
1871+
}
1872+
const result = chatController.truncateRequest(request)
1873+
1874+
// Should keep both images since the first one has 0 chars
1875+
assert.strictEqual(request.conversationState?.currentMessage?.userInputMessage?.images?.length, 2)
1876+
assert.strictEqual(result, 500000 - 400000 - 3.3) // remaining budget after content and second image
1877+
})
1878+
1879+
it('should truncate relevantDocuments and images together with equal priority', () => {
1880+
// 400_000 for content, 100 for doc, 3.3 for image, 100_000 for doc (should be truncated)
1881+
const request: GenerateAssistantResponseCommandInput = {
1882+
conversationState: {
1883+
currentMessage: {
1884+
userInputMessage: {
1885+
content: 'a'.repeat(400_000),
1886+
userInputMessageContext: {
1887+
editorState: {
1888+
relevantDocuments: [
1889+
{ relativeFilePath: 'a', text: 'a'.repeat(100) },
1890+
{ relativeFilePath: 'b', text: 'a'.repeat(100_000) }, // should be truncated
1891+
],
1892+
},
1893+
},
1894+
images: [
1895+
{
1896+
format: 'png',
1897+
source: { bytes: new Uint8Array(1000000000) }, // 3300000 chars
1898+
},
1899+
{
1900+
format: 'png',
1901+
source: { bytes: new Uint8Array(1000) }, // 3.3 chars
1902+
},
1903+
],
1904+
},
1905+
},
1906+
chatTriggerType: undefined,
1907+
},
1908+
}
1909+
const result = chatController.truncateRequest(request)
1910+
// Only the first doc and the image should fit
1911+
assert.strictEqual(
1912+
request.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext?.editorState
1913+
?.relevantDocuments?.length,
1914+
1
1915+
)
1916+
assert.strictEqual(request.conversationState?.currentMessage?.userInputMessage?.images?.length, 1)
1917+
assert.strictEqual(result, 500000 - 400000 - 100 - 3.3)
1918+
})
18061919
})
18071920

18081921
describe('onCreatePrompt', () => {

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

Lines changed: 77 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,11 @@ import {
210210
PaidTierMode,
211211
qProName,
212212
} from '../paidTier/paidTier'
213-
import { Message as DbMessage, messageToStreamingMessage } from './tools/chatDb/util'
213+
import {
214+
estimateCharacterCountFromImageBlock,
215+
Message as DbMessage,
216+
messageToStreamingMessage,
217+
} from './tools/chatDb/util'
214218
import { MODEL_OPTIONS, MODEL_OPTIONS_FOR_REGION } from './constants/modelSelection'
215219
import { DEFAULT_IMAGE_VERIFICATION_OPTIONS, verifyServerImage } from '../../shared/imageVerification'
216220
import { sanitize } from '@aws/lsp-core/out/util/path'
@@ -813,12 +817,6 @@ export class AgenticChatController implements ChatHandlers {
813817
params.context,
814818
params.tabId
815819
)
816-
// Add image context to triggerContext.documentReference for transparency
817-
await this.#additionalContextProvider.appendCustomContextToTriggerContext(
818-
triggerContext,
819-
params.context,
820-
params.tabId
821-
)
822820

823821
let finalResult
824822
if (params.prompt.command === QuickAction.Compact) {
@@ -866,7 +864,7 @@ export class AgenticChatController implements ChatHandlers {
866864
session.conversationId,
867865
token,
868866
triggerContext.documentReference,
869-
additionalContext.filter(item => item.pinned)
867+
additionalContext
870868
)
871869
}
872870

@@ -927,7 +925,7 @@ export class AgenticChatController implements ChatHandlers {
927925
triggerContext: TriggerContext,
928926
additionalContext: AdditionalContentEntryAddition[],
929927
chatResultStream: AgenticChatResultStream,
930-
customContext: ImageBlock[]
928+
images: ImageBlock[]
931929
): Promise<ChatCommandInput> {
932930
this.#debug('Preparing request input')
933931
// Get profileArn from the service manager if available
@@ -943,7 +941,7 @@ export class AgenticChatController implements ChatHandlers {
943941
additionalContext,
944942
session.modelId,
945943
this.#origin,
946-
customContext
944+
images
947945
)
948946
return requestInput
949947
}
@@ -1141,13 +1139,15 @@ export class AgenticChatController implements ChatHandlers {
11411139
conversationIdentifier?: string,
11421140
token?: CancellationToken,
11431141
documentReference?: FileList,
1144-
pinnedContext?: AdditionalContentEntryAddition[]
1142+
additionalContext?: AdditionalContentEntryAddition[]
11451143
): Promise<Result<AgenticChatResultWithMetadata, string>> {
11461144
let currentRequestInput = { ...initialRequestInput }
11471145
let finalResult: Result<AgenticChatResultWithMetadata, string> | null = null
11481146
let iterationCount = 0
11491147
let shouldDisplayMessage = true
11501148
let currentRequestCount = 0
1149+
const pinnedContext = additionalContext?.filter(item => item.pinned)
1150+
11511151
metric.recordStart()
11521152
this.logSystemInformation()
11531153
while (true) {
@@ -1164,7 +1164,7 @@ export class AgenticChatController implements ChatHandlers {
11641164
throw new CancellationError('user')
11651165
}
11661166

1167-
this.truncateRequest(currentRequestInput, pinnedContext)
1167+
this.truncateRequest(currentRequestInput, additionalContext)
11681168
const currentMessage = currentRequestInput.conversationState?.currentMessage
11691169
const conversationId = conversationIdentifier ?? ''
11701170
if (!currentMessage || !conversationId) {
@@ -1248,6 +1248,7 @@ export class AgenticChatController implements ChatHandlers {
12481248
shouldDisplayMessage &&
12491249
!currentMessage.userInputMessage?.content?.startsWith('You are Amazon Q'),
12501250
timestamp: new Date(),
1251+
images: currentMessage.userInputMessage?.images,
12511252
})
12521253
}
12531254
}
@@ -1473,7 +1474,7 @@ export class AgenticChatController implements ChatHandlers {
14731474
* Returns the remaining character budget for chat history.
14741475
* @param request
14751476
*/
1476-
truncateRequest(request: ChatCommandInput, pinnedContext?: AdditionalContentEntryAddition[]): number {
1477+
truncateRequest(request: ChatCommandInput, additionalContext?: AdditionalContentEntryAddition[]): number {
14771478
// TODO: Confirm if this limit applies to SendMessage and rename this constant
14781479
let remainingCharacterBudget = GENERATE_ASSISTANT_RESPONSE_INPUT_LIMIT
14791480
if (!request?.conversationState?.currentMessage?.userInputMessage) {
@@ -1494,22 +1495,69 @@ export class AgenticChatController implements ChatHandlers {
14941495
}
14951496
}
14961497

1497-
// 2. try to fit @context into budget
1498-
let truncatedRelevantDocuments = []
1498+
// 2. try to fit @context and images into budget together
1499+
const docs =
1500+
request.conversationState.currentMessage.userInputMessage.userInputMessageContext?.editorState
1501+
?.relevantDocuments ?? []
1502+
const images = request.conversationState.currentMessage.userInputMessage.images ?? []
1503+
1504+
// Combine docs and images, preserving the order from additionalContext
1505+
let combined
1506+
if (additionalContext && additionalContext.length > 0) {
1507+
let docIdx = 0
1508+
let imageIdx = 0
1509+
combined = additionalContext
1510+
.map(entry => {
1511+
if (entry.type === 'image') {
1512+
return { type: 'image', value: images[imageIdx++] }
1513+
} else {
1514+
return { type: 'doc', value: docs[docIdx++] }
1515+
}
1516+
})
1517+
.filter(item => item.value !== undefined)
1518+
} else {
1519+
combined = [
1520+
...docs.map(d => ({ type: 'doc', value: d })),
1521+
...images.map(i => ({ type: 'image', value: i })),
1522+
]
1523+
}
1524+
1525+
const truncatedDocs: typeof docs = []
1526+
const truncatedImages: typeof images = []
1527+
for (const item of combined) {
1528+
let itemLength = 0
1529+
if (item.type === 'doc') {
1530+
itemLength = (item.value as any)?.text?.length || 0
1531+
if (remainingCharacterBudget >= itemLength) {
1532+
truncatedDocs.push(item.value as (typeof docs)[number])
1533+
remainingCharacterBudget -= itemLength
1534+
}
1535+
} else if (item.type === 'image') {
1536+
// Type guard: only call on ImageBlock
1537+
if (item.value && typeof item.value === 'object' && 'format' in item.value && 'source' in item.value) {
1538+
itemLength = estimateCharacterCountFromImageBlock(item.value)
1539+
if (remainingCharacterBudget >= itemLength) {
1540+
truncatedImages.push(item.value as (typeof images)[number])
1541+
remainingCharacterBudget -= itemLength
1542+
}
1543+
}
1544+
}
1545+
}
1546+
1547+
// Assign truncated lists back to request
14991548
if (
15001549
request.conversationState.currentMessage.userInputMessage.userInputMessageContext?.editorState
15011550
?.relevantDocuments
15021551
) {
1503-
for (const relevantDoc of request.conversationState.currentMessage.userInputMessage.userInputMessageContext
1504-
?.editorState?.relevantDocuments) {
1505-
const docLength = relevantDoc?.text?.length || 0
1506-
if (remainingCharacterBudget > docLength) {
1507-
truncatedRelevantDocuments.push(relevantDoc)
1508-
remainingCharacterBudget = remainingCharacterBudget - docLength
1509-
}
1510-
}
15111552
request.conversationState.currentMessage.userInputMessage.userInputMessageContext.editorState.relevantDocuments =
1512-
truncatedRelevantDocuments
1553+
truncatedDocs
1554+
}
1555+
1556+
if (
1557+
request.conversationState.currentMessage.userInputMessage.images !== undefined &&
1558+
request.conversationState.currentMessage.userInputMessage.images.length > 0
1559+
) {
1560+
request.conversationState.currentMessage.userInputMessage.images = truncatedImages
15131561
}
15141562

15151563
// 3. try to fit current file context
@@ -1528,6 +1576,8 @@ export class AgenticChatController implements ChatHandlers {
15281576
truncatedCurrentDocument
15291577
}
15301578

1579+
const pinnedContext = additionalContext?.filter(item => item.pinned)
1580+
15311581
// 4. try to fit pinned context into budget
15321582
if (pinnedContext && pinnedContext.length > 0) {
15331583
remainingCharacterBudget = this.truncatePinnedContext(remainingCharacterBudget, pinnedContext)
@@ -2901,6 +2951,9 @@ export class AgenticChatController implements ChatHandlers {
29012951
useRelevantDocuments: false,
29022952
}
29032953

2954+
// Clear images to avoid passing them again in follow-up toolUse/toolResult loops, as it is may confuse the model
2955+
updatedRequestInput.conversationState!.currentMessage!.userInputMessage!.images = []
2956+
29042957
for (const toolResult of toolResults) {
29052958
this.#debug(`ToolResult: ${JSON.stringify(toolResult)}`)
29062959
updatedRequestInput.conversationState!.currentMessage!.userInputMessage!.userInputMessageContext!.toolResults.push(

0 commit comments

Comments
 (0)