Skip to content

Commit 5361d63

Browse files
Randall-JiangRandall-Jiang
committed
API refactor
1 parent 8c292c5 commit 5361d63

File tree

4 files changed

+3778
-952
lines changed

4 files changed

+3778
-952
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt‎

Lines changed: 76 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.amazonqCodeTest
55

66
import com.fasterxml.jackson.core.JsonParseException
77
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
8+
import com.fasterxml.jackson.module.kotlin.readValue
89
import com.intellij.openapi.application.ApplicationManager
910
import com.intellij.openapi.components.Service
1011
import com.intellij.openapi.components.service
@@ -34,7 +35,11 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.CodeT
3435
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.Button
3536
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessageContent
3637
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PreviousUTGIterationContext
37-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.ShortAnswer
38+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PackageInfoList
39+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PackageInfo
40+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.TargetFileInfoList
41+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.TargetFileInfo
42+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.CodeReferenceInfo
3843
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.BuildAndExecuteProgressStatus
3944
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
4045
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
@@ -172,7 +177,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
172177
var finished = false
173178
var testGenerationResponse: GetTestGenerationResponse? = null
174179

175-
var shortAnswer = ShortAnswer()
180+
var packageInfoList = PackageInfoList()
176181
LOG.debug {
177182
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
178183
"polling result for id: ${job.testGenerationJobId()}, group name: ${job.testGenerationJobGroupName()}, " +
@@ -187,23 +192,25 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
187192
if (status == TestGenerationJobStatus.COMPLETED) {
188193
LOG.debug {
189194
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
190-
"Test generation completed, short answer string: ${testGenerationResponse.testGenerationJob().shortAnswer()}"
195+
"Test generation completed, package info: ${testGenerationResponse.testGenerationJob().packageInfoList()}"
191196
}
192197
finished = true
193-
if (testGenerationResponse.testGenerationJob().shortAnswer() != null) {
194-
shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
198+
if (testGenerationResponse.testGenerationJob().packageInfoList() != null) {
199+
packageInfoList = parsePackageInfoList(testGenerationResponse.testGenerationJob().packageInfoList())
200+
session.packageInfoList = packageInfoList
195201

196-
val testFileName = shortAnswer.testFilePath?.let { File(it).name }.orEmpty()
202+
val targetFileInfo = packageInfoList.member?.targetFileInfoList?.member?.firstOrNull()
203+
val testFileName = targetFileInfo?.testFilePath?.let { File(it).name }.orEmpty()
197204
session.testFileName = testFileName
198205
// Setting default value to 0 if the value is null or invalid
199-
session.numberOfUnitTestCasesGenerated = shortAnswer.numberOfTestMethods
200-
session.testFileRelativePathToProjectRoot = getTestFilePathRelativeToRoot(shortAnswer)
206+
session.numberOfUnitTestCasesGenerated = targetFileInfo?.numberOfTestMethods ?: 0
207+
session.testFileRelativePathToProjectRoot = getTestFilePathRelativeToRoot(targetFileInfo)
201208

202209
// update test summary card in success case
203210
if (previousIterationContext == null) {
204211
codeTestChatHelper.updateAnswer(
205212
CodeTestChatMessageContent(
206-
message = generateSummaryMessage(path.fileName.toString()) + shortAnswer.planSummary,
213+
message = generateSummaryMessage(path.fileName.toString()) + (packageInfoList.member?.packageSummary ?: ""),
207214
type = ChatMessageType.Answer,
208215
footer = listOf(testFileName)
209216
),
@@ -222,14 +229,17 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
222229
} else if (status == TestGenerationJobStatus.FAILED) {
223230
LOG.debug {
224231
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
225-
"Test generation failed, short answer string: ${testGenerationResponse.testGenerationJob().shortAnswer()}"
226-
}
227-
if (testGenerationResponse.testGenerationJob().shortAnswer() != null) {
228-
shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
229-
if (shortAnswer.stopIteration == "true") {
230-
throw CodeTestException("TestGenFailedError: ${shortAnswer.planSummary}", "TestGenFailedError", shortAnswer.planSummary)
231-
}
232+
"Test generation failed, package info: ${testGenerationResponse.testGenerationJob().packageInfoList()}"
232233
}
234+
// if (testGenerationResponse.testGenerationJob().packageInfoList() != null) {
235+
// shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
236+
// if (shortAnswer.stopIteration == "true") {
237+
// throw CodeTestException("TestGenFailedError: ${shortAnswer.planSummary}", "TestGenFailedError", shortAnswer.planSummary)
238+
// }
239+
// }
240+
241+
// TO DO
242+
// add TestGenerationJobStatus.STOPPED status
233243

234244
// If job status is Failed and has no ShortAnswer then there might be some issue in the backend.
235245
throw CodeTestException(
@@ -245,15 +255,13 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
245255
}
246256
val progressRate = testGenerationResponse.testGenerationJob().progressRate() ?: 0
247257

248-
if (previousIterationContext == null && testGenerationResponse.testGenerationJob().shortAnswer() != null) {
249-
shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
250-
if (shortAnswer.stopIteration == "true") {
251-
throw CodeTestException("TestGenFailedError: ${shortAnswer.planSummary}", "TestGenFailedError", shortAnswer.planSummary)
252-
}
253-
val fileName = shortAnswer.sourceFilePath?.let { Path.of(it).fileName.toString() } ?: path.fileName.toString()
258+
if (previousIterationContext == null && testGenerationResponse.testGenerationJob().packageInfoList() != null) {
259+
packageInfoList = parsePackageInfoList(testGenerationResponse.testGenerationJob().packageInfoList())
260+
val targetFileInfo = packageInfoList.member?.targetFileInfoList?.member?.firstOrNull()
261+
val fileName = targetFileInfo?.filePath?.let { Path.of(it).fileName.toString() }?:path.fileName.toString()
254262
codeTestChatHelper.updateAnswer(
255263
CodeTestChatMessageContent(
256-
message = generateSummaryMessage(fileName) + shortAnswer.planSummary,
264+
message = generateSummaryMessage(fileName) + (packageInfoList.member?.packagePlan ?: ""),
257265
type = ChatMessageType.Answer
258266
),
259267
messageIdOverride = codeTestResponseContext.testSummaryMessageId
@@ -305,14 +313,15 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
305313
return
306314
}
307315

308-
val codeReference = shortAnswer.codeReferences?.map { ref ->
316+
val targetFileInfo = packageInfoList.member?.targetFileInfoList?.member?.firstOrNull()
317+
val codeReference = targetFileInfo?.codeReferences?.map { ref ->
309318
CodeReference(
310319
licenseName = ref.licenseName,
311320
url = ref.url,
312321
information = "${ref.licenseName} - <a href=\"${ref.url}\">${ref.repository}</a>"
313322
)
314323
}
315-
shortAnswer.codeReferences?.let { session.codeReferences = it }
324+
targetFileInfo?.codeReferences?.let { session.codeReferences = it }
316325
val isReferenceAllowed = CodeWhispererSettings.getInstance().isIncludeCodeWithReference()
317326
if (!isReferenceAllowed && codeReference?.isNotEmpty() == true) {
318327
codeTestChatHelper.addAnswer(
@@ -333,7 +342,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
333342
""".trimIndent(),
334343
type = ChatMessageType.Answer,
335344
buttons = listOf(Button("utg_view_diff", "View Diff", keepCardAfterClick = true, position = "outside", status = "info")),
336-
fileList = listOf(getTestFilePathRelativeToRoot(shortAnswer)),
345+
fileList = listOf(getTestFilePathRelativeToRoot(targetFileInfo)),
337346
projectRootName = project.name,
338347
canBeVoted = true,
339348
codeReference = codeReference
@@ -350,7 +359,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
350359
CodeTestChatMessageContent(
351360
type = ChatMessageType.Answer,
352361
buttons = listOf(Button("utg_view_diff", "View Diff", keepCardAfterClick = true, position = "outside", status = "info")),
353-
fileList = listOf(getTestFilePathRelativeToRoot(shortAnswer)),
362+
fileList = listOf(getTestFilePathRelativeToRoot(targetFileInfo)),
354363
projectRootName = project.name,
355364
codeReference = codeReference
356365
),
@@ -374,42 +383,58 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
374383

375384
// Input: test file path relative to project root's parent .
376385
// Output: test file path relative to project root.
377-
// shortAnswer.testFilePath has a format of <projectName>/<test file path relative to project root>.
386+
// targetFileInfo.testFilePath has a format of <projectName>/<test file path relative to project root>.
378387
// test file path in generatedTestDiffs map has a format of resultArtifacts/<test file path relative to project root>.
379388
// both needs to be handled the same way which is remove the first sub-directory
380-
private fun getTestFilePathRelativeToRoot(shortAnswer: ShortAnswer): String {
381-
val pathString = shortAnswer.testFilePath ?: generatedTestDiffs.keys.firstOrNull() ?: throw RuntimeException("No test file path found")
389+
private fun getTestFilePathRelativeToRoot(targetFileInfo: TargetFileInfo?): String {
390+
val pathString = targetFileInfo?.testFilePath ?: generatedTestDiffs.keys.firstOrNull() ?: throw RuntimeException("No test file path found")
382391
val path = Paths.get(pathString)
383392
val updatedPath = path.subpath(1, path.nameCount).toString()
384393
return updatedPath
385394
}
386395

387-
private fun parseShortAnswerString(shortAnswerString: String): ShortAnswer {
388-
// Step 1: Replace single quotes with double quotes
389-
var jsonString = shortAnswerString.replace("'", "\"").replace("```", "")
396+
private fun parsePackageInfoList(packageInfoData: Any): PackageInfoList {
397+
if(packageInfoData is PackageInfoList){
398+
return packageInfoData
399+
}
390400

391-
// Step 2: Replace Python's None with JSON's null
392-
jsonString = jsonString.replace(": None", ": null")
401+
if(packageInfoData is String){
402+
// Step 1: Replace single quotes with double quotes
403+
var jsonString = packageInfoData.replace("'", "\"").replace("```", "")
404+
405+
// Step 2: Replace Python's None with JSON's null
406+
jsonString = jsonString.replace(": None", ": null")
407+
408+
// Step 3: remove extra quotes in the head and tail
409+
if (jsonString.startsWith("\"") && jsonString.endsWith("\"")) {
410+
jsonString = jsonString.substring(1, jsonString.length - 1) // Remove the first and last quote
411+
}
412+
413+
// Step 4: unescape it
414+
jsonString = jsonString.replace("\\\"", "\"")
415+
.replace("\\\\", "\\")
416+
// Deserialize JSON to Kotlin data class
417+
try {
418+
val packageInfoList: PackageInfoList = mapper.readValue(jsonString, PackageInfoList::class.java)
419+
return packageInfoList
420+
} catch (e: JsonParseException) {
421+
LOG.debug(e) { "Test Generation JSON parsing error: ${e.message}" }
422+
throw e
423+
} catch (e: Exception) {
424+
LOG.debug(e) { "Error parsing JSON" }
425+
throw e
426+
}
393427

394-
// Step 3: remove extra quotes in the head and tail
395-
if (jsonString.startsWith("\"") && jsonString.endsWith("\"")) {
396-
jsonString = jsonString.substring(1, jsonString.length - 1) // Remove the first and last quote
397428
}
398429

399-
// Step 4: unescape it
400-
jsonString = jsonString.replace("\\\"", "\"")
401-
.replace("\\\\", "\\")
402-
// Deserialize JSON to Kotlin data class
403-
try {
404-
val shortAnswer: ShortAnswer = mapper.readValue(jsonString, ShortAnswer::class.java)
405-
return shortAnswer
406-
} catch (e: JsonParseException) {
407-
LOG.debug(e) { "Test Generation JSON parsing error: ${e.message}" }
408-
throw e
409-
} catch (e: Exception) {
410-
LOG.debug(e) { "Error parsing JSON" }
430+
try{
431+
return mapper.convertValue(packageInfoData,PackageInfoList::class.java)
432+
}catch (e: JsonParseException) {
433+
LOG.debug(e) { "Failed to covert packageInfoData to PackageInfoList" }
411434
throw e
412435
}
436+
437+
413438
}
414439

415440
private fun storeGeneratedTestDiffs(byteArray: ByteArray, session: Session) {
@@ -507,7 +532,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
507532
private fun resetTestGenFlowSession(session: Session) {
508533
// session.selectedFile doesn't need to be reset since it will remain unchanged
509534
session.conversationState = ConversationState.IN_PROGRESS
510-
session.shortAnswer = ShortAnswer()
535+
session.packageInfoList = PackageInfoList()
511536
session.openedDiffFile = null
512537
session.testFileRelativePathToProjectRoot = ""
513538
session.testFileName = ""

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/model/PackageAndTargetFileInfo.kt‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonqCodeTest.model
5+
6+
import com.fasterxml.jackson.annotation.JsonIgnoreProperties
7+
import com.intellij.testIntegration.TestFramework
8+
9+
data class PackageInfoList(
10+
val member: PackageInfo? = null
11+
)
12+
13+
@JsonIgnoreProperties(ignoreUnknown = true)
14+
data class PackageInfo (
15+
val executionCommand: String? = null,
16+
val buildCommand: String? =null,
17+
val buildOrder: Int? =null,
18+
val testFramework: String? = null,
19+
val packageSummary: String? = null,
20+
val packagePlan: String? = null,
21+
val targetFileInfoList: TargetFileInfoList? = null,
22+
)
23+
24+
data class TargetFileInfoList(
25+
val member: List<TargetFileInfo>? = null,
26+
)
27+
28+
@JsonIgnoreProperties(ignoreUnknown = true)
29+
data class TargetFileInfo(
30+
val filePath: String? = null,
31+
val testFilePath: String? = null,
32+
val testCoverage: Int? = null,
33+
val fileSummary: String? = null,
34+
val filePlan: String? = null,
35+
val codeReferences: List<CodeReferenceInfo>? =null,
36+
val numberOfTestMethods: Int? = null,
37+
)
38+
39+
data class CodeReferenceInfo(
40+
val licenseName: String? = null,
41+
val repository: String? = null,
42+
val url: String? = null,
43+
val recommendationContentSpan: RecommendationContentSpan? = null,
44+
) {
45+
data class RecommendationContentSpan(
46+
val start: Int,
47+
val end: Int,
48+
)
49+
}
50+

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/session/Session.kt‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@ package software.aws.toolkits.jetbrains.services.amazonqCodeTest.session
55

66
import com.intellij.openapi.vfs.VirtualFile
77
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.ConversationState
8-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.ShortAnswer
9-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.ShortAnswerReference
8+
//import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.ShortAnswer
9+
//import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.ShortAnswerReference
10+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PackageInfoList
11+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.CodeReferenceInfo
1012
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
1113
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererUnknownLanguage
1214

@@ -37,14 +39,16 @@ data class Session(val tabId: String) {
3739
// First iteration will have a value of 1
3840
var iteration: Int = 0
3941
var projectRoot: String = "/"
40-
var shortAnswer: ShortAnswer = ShortAnswer()
42+
//var shortAnswer: ShortAnswer = ShortAnswer()
43+
var packageInfoList: PackageInfoList = PackageInfoList()
4144
var selectedFile: VirtualFile? = null
4245
var testFileRelativePathToProjectRoot: String = ""
4346
var testFileName: String = ""
4447
var viewDiffMessageId: String? = null
4548
var openedDiffFile: VirtualFile? = null
4649
val generatedTestDiffs = mutableMapOf<String, String>()
47-
var codeReferences: List<ShortAnswerReference>? = null
50+
//var codeReferences: List<ShortAnswerReference>? = null
51+
var codeReferences: List<CodeReferenceInfo>? = null
4852

4953
// Build loop execution
5054
val buildAndExecuteTaskContext = BuildAndExecuteTaskContext()

0 commit comments

Comments
 (0)