Skip to content

Commit fbdd4d3

Browse files
Will-ShaoHuaWill-ShaoHua
committed
add inline completion project context call and refactor projectContextProvider implementation
1 parent bf33e46 commit fbdd4d3

File tree

3 files changed

+138
-89
lines changed

3 files changed

+138
-89
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/project/InlineBm25Chunk.kt‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.cwc.editor.context.project
5+
6+
data class InlineBm25Chunk(
7+
val content: String,
8+
val filePath: String,
9+
val score: Double,
10+
)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/project/LspMessage.kt‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.cwc.editor.context.project
5+
6+
sealed interface LspMessage {
7+
val endpoint: String
8+
9+
data object Initialize : LspMessage {
10+
override val endpoint: String = "initialize"
11+
}
12+
13+
data object Index : LspMessage {
14+
override val endpoint: String = "buildIndex"
15+
}
16+
17+
data object UpdateIndex : LspMessage {
18+
override val endpoint: String = "updateIndexV2"
19+
}
20+
21+
data object QueryChat : LspMessage {
22+
override val endpoint: String = "query"
23+
}
24+
25+
data object QueryInlineCompletion: LspMessage {
26+
override val endpoint: String = "queryInlineProjectContext"
27+
}
28+
29+
data object GetUsageMetrics : LspMessage {
30+
override val endpoint: String = "getUsage"
31+
}
32+
}
33+
34+
interface LspRequest
35+
36+
data class IndexRequest(
37+
val filePaths: List<String>,
38+
val projectRoot: String,
39+
val config: String,
40+
val language: String = ""
41+
): LspRequest
42+
43+
data class UpdateIndexRequest(
44+
val filePaths: List<String>,
45+
val mode: String,
46+
): LspRequest
47+
48+
data class QueryChatRequest(
49+
val query: String,
50+
): LspRequest
51+
52+
data class QueryInlineCompletionRequest(
53+
val query: String,
54+
val filePath: String
55+
): LspRequest
56+
57+
data class LspResponse(
58+
val responseCode: Int,
59+
val responseBody: String,
60+
)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/project/ProjectContextProvider.kt‎

Lines changed: 68 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -53,25 +53,12 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
5353
}
5454
}
5555
}
56-
data class IndexRequestPayload(
57-
val filePaths: List<String>,
58-
val projectRoot: String,
59-
val refresh: Boolean,
60-
)
6156

6257
data class FileCollectionResult(
6358
val files: List<String>,
6459
val fileSize: Int,
6560
)
6661

67-
data class QueryRequestPayload(
68-
val query: String,
69-
)
70-
71-
data class UpdateIndexRequestPayload(
72-
val filePath: String,
73-
)
74-
7562
data class Usage(
7663
@JsonIgnoreProperties(ignoreUnknown = true)
7764
@JsonProperty("memoryUsage")
@@ -130,37 +117,27 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
130117
}
131118

132119
private fun initEncryption(): Boolean {
133-
logger.info { "project context: init key for ${project.guessProjectDir()} on port ${encoderServer.port}" }
134-
val url = URL("http://localhost:${encoderServer.port}/initialize")
135-
val payload = encoderServer.getEncryptionRequest()
136-
val connection = url.openConnection() as HttpURLConnection
137-
setConnectionProperties(connection)
138-
setConnectionRequest(connection, payload)
139-
logger.info { "project context initialize response code: ${connection.responseCode} for ${project.name}" }
140-
return connection.responseCode == 200
120+
val request = encoderServer.getEncryptionRequest()
121+
val response = sendMsgToLsp(LspMessage.Initialize, request)
122+
return response.responseCode == 200
141123
}
142124

143125
fun index(): Boolean {
144-
logger.info { "project context: indexing ${project.name} on port ${encoderServer.port}" }
126+
val projectRoot = project.guessProjectDir()?.path ?: return false
127+
145128
val indexStartTime = System.currentTimeMillis()
146-
val url = URL("http://localhost:${encoderServer.port}/indexFiles")
147129
val filesResult = collectFiles()
148130
var duration = (System.currentTimeMillis() - indexStartTime).toDouble()
149-
logger.debug { "project context file collection time: ${duration}ms" }
150-
logger.debug { "list of files collected: ${filesResult.files.joinToString("\n")}" }
151-
val projectRoot = project.guessProjectDir()?.path ?: return false
152-
val payload = IndexRequestPayload(filesResult.files, projectRoot, false)
153-
val payloadJson = mapper.writeValueAsString(payload)
154-
val encrypted = encoderServer.encrypt(payloadJson)
155-
156-
val connection = url.openConnection() as HttpURLConnection
157-
setConnectionProperties(connection)
158-
setConnectionRequest(connection, encrypted)
159-
logger.info { "project context index response code: ${connection.responseCode} for ${project.name}" }
131+
logger.debug { "time elapsed to collect project context files: ${duration}ms, collected ${filesResult.files.size} files" }
132+
133+
val encrypted = encryptRequest(IndexRequest(filesResult.files, projectRoot, "all", ""))
134+
val response = sendMsgToLsp(LspMessage.Index, encrypted)
135+
160136
duration = (System.currentTimeMillis() - indexStartTime).toDouble()
161-
val startUrl = getStartUrl(project)
162137
logger.debug { "project context index time: ${duration}ms" }
163-
if (connection.responseCode == 200) {
138+
139+
val startUrl = getStartUrl(project)
140+
if (response.responseCode == 200) {
164141
val usage = getUsage()
165142
TelemetryHelper.recordIndexWorkspace(duration, filesResult.files.size, filesResult.fileSize, true, usage?.memoryUsage, usage?.cpuUsage, startUrl)
166143
logger.debug { "project context index finished for ${project.name}" }
@@ -171,73 +148,47 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
171148
}
172149
}
173150

151+
// TODO: rename queryChat
174152
fun query(prompt: String): List<RelevantDocument> {
175-
logger.info { "project context: querying ${project.name} on port ${encoderServer.port}" }
176-
val url = URL("http://localhost:${encoderServer.port}/query")
177-
val payload = QueryRequestPayload(prompt)
178-
val payloadJson = mapper.writeValueAsString(payload)
179-
val encrypted = encoderServer.encrypt(payloadJson)
180-
181-
val connection = url.openConnection() as HttpURLConnection
182-
setConnectionProperties(connection)
183-
setConnectionTimeout(connection)
184-
setConnectionRequest(connection, encrypted)
185-
186-
val responseCode = connection.responseCode
187-
logger.info { "project context query response code: $responseCode for $prompt" }
188-
val responseBody = if (responseCode == 200) {
189-
connection.inputStream.bufferedReader().use { reader -> reader.readText() }
190-
} else {
191-
""
153+
val encrypted = encryptRequest(QueryChatRequest(prompt))
154+
val response = sendMsgToLsp(LspMessage.QueryChat, encrypted)
155+
156+
return try {
157+
val parsedResponse = mapper.readValue<List<Chunk>>(response.responseBody)
158+
queryResultToRelevantDocuments(parsedResponse)
159+
} catch (e: Exception) {
160+
logger.warn { "error parsing query response ${e.message}" }
161+
emptyList()
192162
}
193-
connection.disconnect()
194-
try {
195-
val parsedResponse = mapper.readValue<List<Chunk>>(responseBody)
196-
return queryResultToRelevantDocuments(parsedResponse)
163+
}
164+
165+
fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> {
166+
val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath))
167+
val response = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted)
168+
169+
return try {
170+
mapper.readValue<List<InlineBm25Chunk>>(response.responseBody)
197171
} catch (e: Exception) {
198172
logger.warn { "error parsing query response ${e.message}" }
199-
return emptyList()
173+
throw e
200174
}
201175
}
202176

203177
private fun getUsage(): Usage? {
204-
logger.info { "project context: getting usage for ${project.name} on port ${encoderServer.port}" }
205-
val url = URL("http://localhost:${encoderServer.port}/getUsage")
206-
val connection = url.openConnection() as HttpURLConnection
207-
setConnectionProperties(connection)
208-
val responseCode = connection.responseCode
209-
210-
logger.info { "project context getUsage response code: $responseCode for ${project.name} " }
211-
val responseBody = if (responseCode == 200) {
212-
connection.inputStream.bufferedReader().use { reader -> reader.readText() }
213-
} else {
214-
""
215-
}
216-
connection.disconnect()
217-
try {
218-
val parsedResponse = mapper.readValue<Usage>(responseBody)
219-
return parsedResponse
178+
val response = sendMsgToLsp(LspMessage.GetUsageMetrics, request = null)
179+
return try {
180+
val parsedResponse = mapper.readValue<Usage>(response.responseBody)
181+
parsedResponse
220182
} catch (e: Exception) {
221183
logger.warn { "error parsing query response ${e.message}" }
222-
return null
184+
null
223185
}
224186
}
225187

226-
fun updateIndex(filePath: String) {
188+
fun updateIndex(filePaths: List<String>, mode: String) {
227189
if (!isIndexComplete.get()) return
228-
logger.info { "project context: updating index for $filePath on port ${encoderServer.port}" }
229-
val url = URL("http://localhost:${encoderServer.port}/updateIndex")
230-
val payload = UpdateIndexRequestPayload(filePath)
231-
val payloadJson = mapper.writeValueAsString(payload)
232-
val encrypted = encoderServer.encrypt(payloadJson)
233-
with(url.openConnection() as HttpURLConnection) {
234-
setConnectionProperties(this)
235-
setConnectionTimeout(this)
236-
setConnectionRequest(this, encrypted)
237-
val responseCode = responseCode
238-
logger.debug { "project context update index response code: $responseCode for $filePath" }
239-
return
240-
}
190+
val encrypted = encryptRequest(UpdateIndexRequest(filePaths, mode))
191+
sendMsgToLsp(LspMessage.UpdateIndex, encrypted)
241192
}
242193

243194
private fun setConnectionTimeout(connection: HttpURLConnection) {
@@ -325,6 +276,34 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
325276
return documents
326277
}
327278

279+
private fun encryptRequest(r: LspRequest): String {
280+
val payloadJson = mapper.writeValueAsString(r)
281+
return encoderServer.encrypt(payloadJson)
282+
}
283+
284+
private fun sendMsgToLsp(msgType: LspMessage, request: String?): LspResponse {
285+
logger.info { "sending message: ${msgType.endpoint} to lsp on port ${encoderServer.port}" }
286+
val url = URL("http://localhost:${encoderServer.port}/${msgType.endpoint}")
287+
288+
return with(url.openConnection() as HttpURLConnection) {
289+
setConnectionProperties(this)
290+
setConnectionTimeout(this)
291+
request?.let { r ->
292+
setConnectionRequest(this, r)
293+
}
294+
val responseCode = this.responseCode
295+
logger.info { "receiving response for $msgType with responseCode $responseCode" }
296+
297+
val responseBody = if (responseCode == 200) {
298+
this.inputStream.bufferedReader().use { reader -> reader.readText() }
299+
} else {
300+
""
301+
}
302+
303+
LspResponse(responseCode, responseBody)
304+
}
305+
}
306+
328307
override fun dispose() {
329308
retryCount.set(0)
330309
}

0 commit comments

Comments
 (0)