Skip to content

Commit 6af6a46

Browse files
Will-ShaoHuaWill-ShaoHua
authored
refactor(amazonq): move codewhisperer utils/contants to amazonq/shared (#4970)
1 parent 4319d6d commit 6af6a46

File tree

11 files changed

+128
-83
lines changed

11 files changed

+128
-83
lines changed

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt‎

Lines changed: 12 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import com.intellij.openapi.Disposable
77
import com.intellij.openapi.application.ApplicationManager
88
import com.intellij.openapi.components.service
99
import com.intellij.openapi.project.Project
10-
import com.intellij.openapi.util.SystemInfo
1110
import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider
1211
import software.amazon.awssdk.services.codewhisperer.CodeWhispererClient
1312
import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanRequest
@@ -25,13 +24,10 @@ import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUr
2524
import software.amazon.awssdk.services.codewhispererruntime.model.Dimension
2625
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
2726
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
28-
import software.amazon.awssdk.services.codewhispererruntime.model.IdeCategory
2927
import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest
3028
import software.amazon.awssdk.services.codewhispererruntime.model.ListFeatureEvaluationsResponse
31-
import software.amazon.awssdk.services.codewhispererruntime.model.OperatingSystem
3229
import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventResponse
3330
import software.amazon.awssdk.services.codewhispererruntime.model.SuggestionState
34-
import software.amazon.awssdk.services.codewhispererruntime.model.UserContext
3531
import software.amazon.awssdk.services.codewhispererruntime.model.UserIntent
3632
import software.aws.toolkits.core.utils.debug
3733
import software.aws.toolkits.core.utils.getLogger
@@ -43,6 +39,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
4339
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
4440
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
4541
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
42+
import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext
4643
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
4744
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
4845
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
@@ -51,10 +48,8 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestCon
5148
import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContextNew
5249
import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
5350
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
54-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.FEATURE_EVALUATION_PRODUCT_NAME
5551
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
5652
import software.aws.toolkits.jetbrains.services.codewhisperer.util.transform
57-
import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
5853
import software.aws.toolkits.telemetry.CodewhispererCompletionType
5954
import software.aws.toolkits.telemetry.CodewhispererSuggestionState
6055
import java.time.Instant
@@ -201,24 +196,6 @@ interface CodeWhispererClientAdaptor : Disposable {
201196
}
202197

203198
open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispererClientAdaptor {
204-
private val codeWhispererUserContext = ClientMetadata.getDefault().let {
205-
val osForCodeWhisperer: OperatingSystem =
206-
when {
207-
SystemInfo.isWindows -> OperatingSystem.WINDOWS
208-
SystemInfo.isMac -> OperatingSystem.MAC
209-
// For now, categorize everything else as "Linux" (Linux/FreeBSD/Solaris/etc)
210-
else -> OperatingSystem.LINUX
211-
}
212-
213-
UserContext.builder()
214-
.ideCategory(IdeCategory.JETBRAINS)
215-
.operatingSystem(osForCodeWhisperer)
216-
.product(FEATURE_EVALUATION_PRODUCT_NAME)
217-
.clientId(it.clientId)
218-
.ideVersion(it.awsVersion)
219-
.build()
220-
}
221-
222199
private val mySigv4Client by lazy { createUnmanagedSigv4Client() }
223200

224201
@Volatile
@@ -348,7 +325,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
348325
}
349326
}
350327
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
351-
requestBuilder.userContext(codeWhispererUserContext)
328+
requestBuilder.userContext(codeWhispererUserContext())
352329
}
353330
}
354331

@@ -391,7 +368,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
391368
}
392369
}
393370
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
394-
requestBuilder.userContext(codeWhispererUserContext)
371+
requestBuilder.userContext(codeWhispererUserContext())
395372
}
396373
}
397374

@@ -413,7 +390,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
413390
}
414391
}
415392
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
416-
requestBuilder.userContext(codeWhispererUserContext)
393+
requestBuilder.userContext(codeWhispererUserContext())
417394
}
418395

419396
override fun sendUserModificationTelemetry(
@@ -440,7 +417,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
440417
}
441418
}
442419
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
443-
requestBuilder.userContext(codeWhispererUserContext)
420+
requestBuilder.userContext(codeWhispererUserContext())
444421
}
445422

446423
override fun sendCodeScanTelemetry(
@@ -459,7 +436,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
459436
}
460437
}
461438
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
462-
requestBuilder.userContext(codeWhispererUserContext)
439+
requestBuilder.userContext(codeWhispererUserContext())
463440
}
464441
override fun sendCodeScanRemediationTelemetry(
465442
language: CodeWhispererProgrammingLanguage?,
@@ -489,11 +466,11 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
489466
}
490467
}
491468
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
492-
requestBuilder.userContext(codeWhispererUserContext)
469+
requestBuilder.userContext(codeWhispererUserContext())
493470
}
494471

495472
override fun listFeatureEvaluations(): ListFeatureEvaluationsResponse = bearerClient().listFeatureEvaluations {
496-
it.userContext(codeWhispererUserContext)
473+
it.userContext(codeWhispererUserContext())
497474
}
498475

499476
override fun sendMetricDataTelemetry(eventName: String, metadata: Map<String, Any?>): SendTelemetryEventResponse =
@@ -507,7 +484,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
507484
}
508485
}
509486
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
510-
requestBuilder.userContext(codeWhispererUserContext)
487+
requestBuilder.userContext(codeWhispererUserContext())
511488
}
512489

513490
override fun sendChatAddMessageTelemetry(
@@ -547,7 +524,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
547524
}
548525
}
549526
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
550-
requestBuilder.userContext(codeWhispererUserContext)
527+
requestBuilder.userContext(codeWhispererUserContext())
551528
}
552529

553530
override fun sendChatInteractWithMessageTelemetry(
@@ -576,7 +553,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
576553
telemetryEventBuilder.chatInteractWithMessageEvent(event)
577554
}
578555
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
579-
requestBuilder.userContext(codeWhispererUserContext)
556+
requestBuilder.userContext(codeWhispererUserContext())
580557
}
581558

582559
override fun sendChatUserModificationTelemetry(
@@ -602,7 +579,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
602579
}
603580
}
604581
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
605-
requestBuilder.userContext(codeWhispererUserContext)
582+
requestBuilder.userContext(codeWhispererUserContext())
606583
}
607584

608585
override fun dispose() {

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ import com.intellij.util.xmlb.annotations.Property
2020
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
2121
import software.aws.toolkits.core.utils.debug
2222
import software.aws.toolkits.core.utils.getLogger
23+
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
2324
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
2425
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService
2526
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
26-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfIamIdentityCenterConnection
2727
import software.aws.toolkits.jetbrains.utils.notifyInfo
2828
import software.aws.toolkits.jetbrains.utils.notifyWarn
2929
import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt‎

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,15 @@ import com.intellij.openapi.components.Service
77
import com.intellij.openapi.components.service
88
import com.intellij.openapi.project.Project
99
import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
10+
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
1011
import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue
12+
import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest
1113
import software.aws.toolkits.core.utils.debug
1214
import software.aws.toolkits.core.utils.getLogger
13-
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
14-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfBIDConnection
15-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfIamIdentityCenterConnection
15+
import software.aws.toolkits.jetbrains.core.awsClient
16+
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfBIDConnection
17+
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
18+
import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext
1619
import software.aws.toolkits.jetbrains.utils.isQExpired
1720

1821
@Service
@@ -25,7 +28,9 @@ class CodeWhispererFeatureConfigService {
2528

2629
LOG.debug { "Fetching feature configs" }
2730
try {
28-
val response = CodeWhispererClientAdaptor.getInstance(project).listFeatureEvaluations()
31+
val response = project.awsClient<CodeWhispererRuntimeClient>().listFeatureEvaluations {
32+
it.userContext(codeWhispererUserContext())
33+
}
2934

3035
// Simply force overwrite feature configs from server response, no needed to check existing values.
3136
response.featureEvaluations().forEach {
@@ -49,7 +54,16 @@ class CodeWhispererFeatureConfigService {
4954
val availableCustomizations =
5055
calculateIfIamIdentityCenterConnection(project) {
5156
try {
52-
CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations().map { c -> c.arn }
57+
project.awsClient<CodeWhispererRuntimeClient>().listAvailableCustomizationsPaginator(
58+
ListAvailableCustomizationsRequest.builder().build()
59+
)
60+
.stream()
61+
.toList()
62+
.flatMap { resp ->
63+
resp.customizations().map {
64+
it.arn()
65+
}
66+
}
5367
} catch (e: Exception) {
5468
LOG.debug(e) { "Failed to list available customizations" }
5569
null
@@ -71,9 +85,11 @@ class CodeWhispererFeatureConfigService {
7185
}
7286

7387
fun getFeatureConfigsTelemetry(): String =
74-
"{${featureConfigs.entries.joinToString(", ") { (name, context) ->
75-
"$name: ${context.variation}"
76-
}}}"
88+
"{${
89+
featureConfigs.entries.joinToString(", ") { (name, context) ->
90+
"$name: ${context.variation}"
91+
}
92+
}}"
7793

7894
// TODO: for all feature variations, define a contract that can be enforced upon the implementation of
7995
// the business logic.
@@ -88,7 +104,8 @@ class CodeWhispererFeatureConfigService {
88104
// 6) Add a test case for this feature.
89105
fun getTestFeature(): String = getFeatureValueForKey(TEST_FEATURE_NAME).stringValue()
90106

91-
fun getIsDataCollectionEnabled(): Boolean = getFeatureValueForKey(DATA_COLLECTION_FEATURE).stringValue() == "data-collection"
107+
fun getIsDataCollectionEnabled(): Boolean =
108+
getFeatureValueForKey(DATA_COLLECTION_FEATURE).stringValue() == "data-collection"
92109

93110
fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue()
94111

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import kotlinx.coroutines.delay
1010
import kotlinx.coroutines.isActive
1111
import kotlinx.coroutines.launch
1212
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
13+
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
1314
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
1415
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
1516
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
@@ -20,7 +21,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe
2021
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
2122
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.FEATURE_CONFIG_POLL_INTERVAL_IN_MS
2223
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
23-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfIamIdentityCenterConnection
2424
import software.aws.toolkits.jetbrains.utils.isQConnected
2525
import software.aws.toolkits.jetbrains.utils.isQExpired
2626
import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ object CodeWhispererConstants {
3131
const val POPUP_DELAY_CHECK_INTERVAL: Long = 25
3232
const val IDLE_TIME_CHECK_INTERVAL: Long = 25
3333
const val SUPPLEMENTAL_CONTEXT_TIMEOUT = 50L
34-
const val FEATURE_EVALUATION_PRODUCT_NAME = "CodeWhisperer"
3534

3635
val AWSTemplateKeyWordsRegex = Regex("(AWSTemplateFormatVersion|Resources|AWS::|Description)")
3736
val AWSTemplateCaseInsensitiveKeyWordsRegex = Regex("(cloudformation|cfn|template|description)")

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt‎

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -50,27 +50,6 @@ import software.aws.toolkits.resources.message
5050
import software.aws.toolkits.telemetry.CodewhispererCompletionType
5151
import software.aws.toolkits.telemetry.CodewhispererGettingStartedTask
5252

53-
fun <T> calculateIfIamIdentityCenterConnection(project: Project, calculationTask: (connection: ToolkitConnection) -> T): T? =
54-
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let {
55-
calculateIfIamIdentityCenterConnection(it, calculationTask)
56-
}
57-
58-
fun <T> calculateIfIamIdentityCenterConnection(connection: ToolkitConnection, calculationTask: (connection: ToolkitConnection) -> T): T? =
59-
if (connection.isSono()) {
60-
null
61-
} else {
62-
calculationTask(connection)
63-
}
64-
65-
fun <T> calculateIfBIDConnection(project: Project, calculationTask: (connection: ToolkitConnection) -> T): T? =
66-
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let {
67-
if (it.isSono()) {
68-
calculationTask(it)
69-
} else {
70-
null
71-
}
72-
}
73-
7453
// Controls the condition to send telemetry event to CodeWhisperer service, currently:
7554
// 1. It will be sent for Builder ID users, only if they have optin telemetry sharing.
7655
// 2. It will be sent for IdC users, regardless of telemetry optout status.

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRul
7373
import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule
7474
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
7575
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION
76+
import software.aws.toolkits.jetbrains.services.amazonq.FEATURE_EVALUATION_PRODUCT_NAME
7677
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
7778
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonRequest
7879
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithToken
@@ -87,7 +88,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe
8788
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
8889
import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
8990
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
90-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.FEATURE_EVALUATION_PRODUCT_NAME
9191
import software.aws.toolkits.jetbrains.settings.AwsSettings
9292
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
9393
import software.aws.toolkits.telemetry.CodewhispererCompletionType

0 commit comments

Comments
 (0)