Skip to content

Commit 392a31d

Browse files
nborges-awsnborges-aws
authored
feat: support for project context in Q Chat (#1061) (#1101)
1 parent 8875ddf commit 392a31d

File tree

3 files changed

+287
-1
lines changed

3 files changed

+287
-1
lines changed

‎server/aws-lsp-codewhisperer/src/language-server/chat/contexts/triggerContext.ts‎

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,46 @@
11
import { TriggerType } from '@aws/chat-client-ui-types'
2-
import { ChatTriggerType, UserIntent, Tool, ToolResult } from '@amzn/codewhisperer-streaming'
2+
import { ChatTriggerType, UserIntent, Tool, ToolResult, RelevantTextDocument } from '@amzn/codewhisperer-streaming'
33
import { BedrockTools, ChatParams, CursorState, InlineChatParams } from '@aws/language-server-runtimes/server-interface'
44
import { Features } from '../../types'
55
import { DocumentContext, DocumentContextExtractor } from './documentContext'
66
import { SendMessageCommandInput } from '../../../shared/streamingClientService'
7+
import { LocalProjectContextController } from '../../../shared/localProjectContextController'
8+
import { convertChunksToRelevantTextDocuments } from '../tools/relevantTextDocuments'
79

810
export interface TriggerContext extends Partial<DocumentContext> {
911
userIntent?: UserIntent
1012
triggerType?: TriggerType
13+
useRelevantDocuments?: boolean
14+
relevantDocuments?: RelevantTextDocument[]
1115
}
1216

1317
export class QChatTriggerContext {
1418
private static readonly DEFAULT_CURSOR_STATE: CursorState = { position: { line: 0, character: 0 } }
1519

1620
#workspace: Features['workspace']
1721
#documentContextExtractor: DocumentContextExtractor
22+
#logger: Features['logging']
1823

1924
constructor(workspace: Features['workspace'], logger: Features['logging']) {
2025
this.#workspace = workspace
2126
this.#documentContextExtractor = new DocumentContextExtractor({ logger, workspace })
27+
this.#logger = logger
2228
}
2329

2430
async getNewTriggerContext(params: ChatParams | InlineChatParams): Promise<TriggerContext> {
2531
const documentContext: DocumentContext | undefined = await this.extractDocumentContext(params)
2632

33+
const useRelevantDocuments =
34+
'context' in params
35+
? params.context?.some(context => typeof context !== 'string' && context.command === '@workspace')
36+
: false
37+
const relevantDocuments = useRelevantDocuments ? await this.extractProjectContext(params.prompt.prompt) : []
38+
2739
return {
2840
...documentContext,
2941
userIntent: this.#guessIntentFromPrompt(params.prompt.prompt),
42+
useRelevantDocuments,
43+
relevantDocuments,
3044
}
3145
}
3246

@@ -56,11 +70,21 @@ export class QChatTriggerContext {
5670
programmingLanguage: triggerContext.programmingLanguage,
5771
relativeFilePath: triggerContext.relativeFilePath,
5872
},
73+
...(triggerContext.useRelevantDocuments && {
74+
useRelevantDocuments: triggerContext.useRelevantDocuments,
75+
relevantDocuments: triggerContext.relevantDocuments,
76+
}),
5977
},
6078
tools,
6179
}
6280
: {
6381
tools,
82+
...(triggerContext.useRelevantDocuments && {
83+
editorState: {
84+
useRelevantDocuments: triggerContext.useRelevantDocuments,
85+
relevantDocuments: triggerContext.relevantDocuments,
86+
},
87+
}),
6488
},
6589
userIntent: triggerContext.userIntent,
6690
origin: 'IDE',
@@ -93,6 +117,19 @@ export class QChatTriggerContext {
93117
: undefined
94118
}
95119

120+
async extractProjectContext(query?: string): Promise<RelevantTextDocument[]> {
121+
if (query) {
122+
try {
123+
const contextController = await LocalProjectContextController.getInstance()
124+
const resp = await contextController.queryVectorIndex({ query })
125+
return convertChunksToRelevantTextDocuments(resp)
126+
} catch (e) {
127+
this.#logger.error(`Failed to extract project context for chat trigger: ${e}`)
128+
}
129+
}
130+
return []
131+
}
132+
96133
#guessIntentFromPrompt(prompt?: string): UserIntent | undefined {
97134
if (prompt === undefined) {
98135
return undefined

‎server/aws-lsp-codewhisperer/src/language-server/chat/tools/relevantTextDocuments.test.ts‎

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
import { convertChunksToRelevantTextDocuments } from './relevantTextDocuments'
2+
import { Chunk } from 'local-indexing'
3+
import { RelevantTextDocument } from '@amzn/codewhisperer-streaming'
4+
import * as assert from 'assert'
5+
6+
describe('relevantTextDocuments', () => {
7+
it('converts empty array to empty array', () => {
8+
const result = convertChunksToRelevantTextDocuments([])
9+
assert.deepStrictEqual(result, [])
10+
})
11+
12+
it('combines chunks from same file and sorts by startLine', () => {
13+
const chunks: Chunk[] = [
14+
{
15+
id: '1',
16+
index: 0,
17+
filePath: 'test.ts',
18+
relativePath: 'src/test.ts',
19+
content: 'second chunk',
20+
startLine: 2,
21+
programmingLanguage: 'typescript',
22+
vec: [],
23+
},
24+
{
25+
id: '2',
26+
index: 1,
27+
filePath: 'test.ts',
28+
relativePath: 'src/test.ts',
29+
content: 'first chunk',
30+
startLine: 1,
31+
programmingLanguage: 'typescript',
32+
vec: [],
33+
},
34+
]
35+
36+
const expected: RelevantTextDocument[] = [
37+
{
38+
relativeFilePath: 'src/test.ts',
39+
programmingLanguage: { languageName: 'typescript' },
40+
text: 'first chunk\nsecond chunk',
41+
},
42+
]
43+
44+
const result = convertChunksToRelevantTextDocuments(chunks)
45+
assert.deepStrictEqual(result, expected)
46+
})
47+
48+
it('handles chunks without startLine', () => {
49+
const chunks: Chunk[] = [
50+
{
51+
id: '1',
52+
index: 0,
53+
filePath: 'test.ts',
54+
relativePath: 'src/test.ts',
55+
content: 'chunk1',
56+
programmingLanguage: 'typescript',
57+
vec: [],
58+
},
59+
{
60+
id: '2',
61+
index: 1,
62+
filePath: 'test.ts',
63+
relativePath: 'src/test.ts',
64+
content: 'chunk2',
65+
programmingLanguage: 'typescript',
66+
vec: [],
67+
},
68+
]
69+
70+
const expected: RelevantTextDocument[] = [
71+
{
72+
relativeFilePath: 'src/test.ts',
73+
programmingLanguage: { languageName: 'typescript' },
74+
text: 'chunk1\nchunk2',
75+
},
76+
]
77+
78+
const result = convertChunksToRelevantTextDocuments(chunks)
79+
assert.deepStrictEqual(result, expected)
80+
})
81+
82+
it('handles unknown programming language', () => {
83+
const chunks: Chunk[] = [
84+
{
85+
id: '1',
86+
index: 0,
87+
filePath: 'test.txt',
88+
relativePath: 'src/test.txt',
89+
content: 'content',
90+
programmingLanguage: 'unknown',
91+
vec: [],
92+
},
93+
]
94+
95+
const expected: RelevantTextDocument[] = [
96+
{
97+
relativeFilePath: 'src/test.txt',
98+
text: 'content',
99+
},
100+
]
101+
102+
const result = convertChunksToRelevantTextDocuments(chunks)
103+
assert.deepStrictEqual(result, expected)
104+
})
105+
106+
it('filters out empty content', () => {
107+
const chunks: Chunk[] = [
108+
{
109+
id: '1',
110+
index: 0,
111+
filePath: 'test.ts',
112+
relativePath: 'src/test.ts',
113+
content: '',
114+
programmingLanguage: 'typescript',
115+
vec: [],
116+
},
117+
{
118+
id: '2',
119+
index: 1,
120+
filePath: 'test.ts',
121+
relativePath: 'src/test.ts',
122+
content: 'valid content',
123+
programmingLanguage: 'typescript',
124+
vec: [],
125+
},
126+
]
127+
128+
const expected: RelevantTextDocument[] = [
129+
{
130+
relativeFilePath: 'src/test.ts',
131+
programmingLanguage: { languageName: 'typescript' },
132+
text: 'valid content',
133+
},
134+
]
135+
136+
const result = convertChunksToRelevantTextDocuments(chunks)
137+
assert.deepStrictEqual(result, expected)
138+
})
139+
140+
it('truncates relative file path if too long', () => {
141+
const longPath = 'a'.repeat(5000)
142+
const chunks: Chunk[] = [
143+
{
144+
id: '1',
145+
index: 0,
146+
filePath: 'test.ts',
147+
relativePath: longPath,
148+
content: 'content',
149+
programmingLanguage: 'typescript',
150+
vec: [],
151+
},
152+
]
153+
154+
const result = convertChunksToRelevantTextDocuments(chunks)
155+
assert.strictEqual(result[0].relativeFilePath?.length, 4000)
156+
})
157+
158+
it('handles multiple files', () => {
159+
const chunks: Chunk[] = [
160+
{
161+
id: '1',
162+
index: 0,
163+
filePath: 'test1.ts',
164+
relativePath: 'src/test1.ts',
165+
content: 'content1',
166+
programmingLanguage: 'typescript',
167+
vec: [],
168+
},
169+
{
170+
id: '2',
171+
index: 1,
172+
filePath: 'test2.ts',
173+
relativePath: 'src/test2.ts',
174+
content: 'content2',
175+
programmingLanguage: 'typescript',
176+
vec: [],
177+
},
178+
]
179+
180+
const expected: RelevantTextDocument[] = [
181+
{
182+
relativeFilePath: 'src/test1.ts',
183+
programmingLanguage: { languageName: 'typescript' },
184+
text: 'content1',
185+
},
186+
{
187+
relativeFilePath: 'src/test2.ts',
188+
programmingLanguage: { languageName: 'typescript' },
189+
text: 'content2',
190+
},
191+
]
192+
193+
const result = convertChunksToRelevantTextDocuments(chunks)
194+
assert.deepStrictEqual(result, expected)
195+
})
196+
})

‎server/aws-lsp-codewhisperer/src/language-server/chat/tools/relevantTextDocuments.ts‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import { RelevantTextDocument } from '@amzn/codewhisperer-streaming'
2+
import { Chunk } from 'local-indexing'
3+
4+
export function convertChunksToRelevantTextDocuments(chunks: Chunk[]): RelevantTextDocument[] {
5+
const filePathSizeLimit = 4_000
6+
7+
const groupedChunks = chunks.reduce(
8+
(acc, chunk) => {
9+
const key = chunk.filePath
10+
if (!acc[key]) {
11+
acc[key] = []
12+
}
13+
acc[key].push(chunk)
14+
return acc
15+
},
16+
{} as Record<string, Chunk[]>
17+
)
18+
19+
return Object.entries(groupedChunks).map(([filePath, fileChunks]) => {
20+
fileChunks.sort((a, b) => {
21+
if (a.startLine !== undefined && b.startLine !== undefined) {
22+
return a.startLine - b.startLine
23+
}
24+
return 0
25+
})
26+
27+
const firstChunk = fileChunks[0]
28+
29+
let programmingLanguage
30+
if (firstChunk.programmingLanguage && firstChunk.programmingLanguage !== 'unknown') {
31+
programmingLanguage = {
32+
languageName: firstChunk.programmingLanguage,
33+
}
34+
}
35+
36+
const combinedContent = fileChunks
37+
.map(chunk => chunk.content)
38+
.filter(content => content !== undefined && content !== '')
39+
.join('\n')
40+
41+
const relevantTextDocument: RelevantTextDocument = {
42+
relativeFilePath: firstChunk.relativePath
43+
? firstChunk.relativePath.substring(0, filePathSizeLimit)
44+
: undefined,
45+
programmingLanguage,
46+
text: combinedContent || undefined,
47+
}
48+
49+
return Object.fromEntries(
50+
Object.entries(relevantTextDocument).filter(([_, value]) => value !== undefined)
51+
) as RelevantTextDocument
52+
})
53+
}

0 commit comments

Comments
 (0)