Add logic that uses IAM Q Dev Client

Author: ahusseinali

package-lock.json

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
 import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { ChatSession } from 'aws-core-vscode/codewhispererChat'
@@ -12,7 +13,11 @@ export default {
     chatApi: {
         async chat(request: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput> {
             const chatSession = new ChatSession()
-            return chatSession.chat(request)
+            return chatSession.chatSso(request)
+        },
+        async chatIam(request: SendMessageRequest): Promise<SendMessageCommandOutput> {
+            const chatSession = new ChatSession()
+            return chatSession.chatIam(request)
         },
     },
     authApi: {

packages/amazonq/src/api.ts

@@ -5,6 +5,7 @@
 
 import vscode from 'vscode'
 import { VSCODE_EXTENSION_ID } from '../shared/utilities'
+import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
 import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
 import { FeatureAuthState } from '../codewhisperer/util/authUtil'
 import { ToolkitError } from '../shared'
@@ -16,6 +17,7 @@ import { ToolkitError } from '../shared'
 export interface api {
     chatApi: {
         chat(request: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput>
+        chatIam(request: SendMessageRequest): Promise<SendMessageCommandOutput>
     }
     authApi: {
         reauthIfNeeded(): Promise<void>

packages/core/src/amazonq/extApi.ts

@@ -8,7 +8,7 @@ import * as localizedText from '../../shared/localizedText'
 import { Auth } from '../../auth/auth'
 import { ToolkitError, isNetworkError, tryRun } from '../../shared/errors'
 import { getSecondaryAuth, setScopes } from '../../auth/secondaryAuth'
-import { isCloud9, isSageMaker } from '../../shared/extensionUtilities'
+import { isAmazonQ, isCloud9, isSageMaker } from '../../shared/extensionUtilities'
 import { AmazonQPromptSettings } from '../../shared/settings'
 import {
     scopesCodeWhispererCore,
@@ -60,7 +60,7 @@ export const isValidCodeWhispererCoreConnection = (conn?: Connection): conn is C
     }
 
     if (isSageMaker()) {
-        return isIamConnection(conn)
+        return isIamConnection(conn) || (isSsoConnection(conn) && hasScopes(conn, codeWhispererCoreScopes))
     }
 
     return (
@@ -70,6 +70,10 @@ export const isValidCodeWhispererCoreConnection = (conn?: Connection): conn is C
 }
 /** Superset that includes all of CodeWhisperer + Amazon Q */
 export const isValidAmazonQConnection = (conn?: Connection): conn is Connection => {
+    if (isSageMaker() && isAmazonQ()) {
+        return isIamConnection(conn) || (isSsoConnection(conn) && hasScopes(conn, amazonQScopes))
+    }
+
     return (
         (isSsoConnection(conn) || isBuilderIdConnection(conn)) &&
         isValidCodeWhispererCoreConnection(conn) &&
@@ -442,9 +446,6 @@ export class AuthUtil {
         if (conn === undefined) {
             return buildFeatureAuthState(AuthStates.disconnected)
         }
-        if (!isSsoConnection(conn)) {
-            throw new ToolkitError(`Connection "${conn.id}" is not a valid type: ${conn.type}`)
-        }
 
         // default to expired to indicate reauth is needed if unmodified
         const state: FeatureAuthState = buildFeatureAuthState(AuthStates.expired)
@@ -453,7 +454,7 @@ export class AuthUtil {
             return state
         }
 
-        if (isBuilderIdConnection(conn) || isIdcSsoConnection(conn)) {
+        if (isBuilderIdConnection(conn) || isIdcSsoConnection(conn) || isSageMaker()) {
             if (isValidCodeWhispererCoreConnection(conn)) {
                 state[Features.codewhispererCore] = AuthStates.connected
             }

packages/core/src/codewhisperer/util/authUtil.ts

@@ -3,10 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { SendMessageCommandOutput, SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
 import { GenerateAssistantResponseCommandOutput, GenerateAssistantResponseRequest } from '@amzn/codewhisperer-streaming'
 import * as vscode from 'vscode'
 import { ToolkitError } from '../../../../shared/errors'
 import { createCodeWhispererChatStreamingClient } from '../../../../shared/clients/codewhispererChatClient'
+import { createQDeveloperStreamingClient } from '../../../../shared/clients/qDeveloperChatClient'
 
 export class ChatSession {
     private sessionId?: string
@@ -28,8 +30,28 @@ export class ChatSession {
     public setSessionID(id?: string) {
         this.sessionId = id
     }
+    async chatIam(chatRequest: SendMessageRequest): Promise<SendMessageCommandOutput> {
+        const client = await createQDeveloperStreamingClient()
 
-    async chat(chatRequest: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput> {
+        const response = await client.sendMessage(chatRequest)
+        if (!response.sendMessageResponse) {
+            throw new ToolkitError(
+                `Empty chat response. Session id: ${this.sessionId} Request ID: ${response.$metadata.requestId}`
+            )
+        }
+
+        const responseStream = response.sendMessageResponse
+        for await (const event of responseStream) {
+            if ('messageMetadataEvent' in event) {
+                this.sessionId = event.messageMetadataEvent?.conversationId
+                break
+            }
+        }
+
+        return response
+    }
+
+    async chatSso(chatRequest: GenerateAssistantResponseRequest): Promise<GenerateAssistantResponseCommandOutput> {
         const client = await createCodeWhispererChatStreamingClient()
 
         if (this.sessionId !== undefined && chatRequest.conversationState !== undefined) {

packages/core/src/codewhispererChat/clients/chat/v0/chat.ts

@@ -4,9 +4,9 @@
  */
 
 import {
+    ConversationState,
     CursorState,
     DocumentSymbol,
-    GenerateAssistantResponseRequest,
     RelevantTextDocument,
     SymbolType,
     TextDocument,
@@ -37,7 +37,7 @@ export const supportedLanguagesList = [
 const filePathSizeLimit = 4_000
 const customerMessageSizeLimit = 4_000
 
-export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): GenerateAssistantResponseRequest {
+export function triggerPayloadToChatRequest(triggerPayload: TriggerPayload): { conversationState: ConversationState } {
     let document: TextDocument | undefined = undefined
     let cursorState: CursorState | undefined = undefined
 

packages/core/src/codewhispererChat/controllers/chat/chatRequest/converter.ts

@@ -5,7 +5,7 @@
 import { Event as VSCodeEvent, Uri } from 'vscode'
 import { EditorContextExtractor } from '../../editor/context/extractor'
 import { ChatSessionStorage } from '../../storages/chatSession'
-import { Messenger, StaticTextResponseType } from './messenger/messenger'
+import { Messenger, MessengerResponseType, StaticTextResponseType } from './messenger/messenger'
 import {
     PromptMessage,
     ChatTriggerType,
@@ -34,10 +34,8 @@ import { EditorContentController } from '../../../amazonq/commons/controllers/co
 import { EditorContextCommand } from '../../commands/registerCommands'
 import { PromptsGenerator } from './prompts/promptsGenerator'
 import { TriggerEventsStorage } from '../../storages/triggerEvents'
-import {
-    CodeWhispererStreamingServiceException,
-    GenerateAssistantResponseCommandOutput,
-} from '@amzn/codewhisperer-streaming'
+import { SendMessageRequest } from '@amzn/amazon-q-developer-streaming-client'
+import { CodeWhispererStreamingServiceException, __MetadataBearer } from '@amzn/codewhisperer-streaming'
 import { UserIntentRecognizer } from './userIntent/userIntentRecognizer'
 import { CWCTelemetryHelper, recordTelemetryChatRunCommand } from './telemetryHelper'
 import { CodeWhispererTracker } from '../../../codewhisperer/tracker/codewhispererTracker'
@@ -54,6 +52,7 @@ import { getHttpStatusCode, AwsClientResponseError } from '../../../shared/error
 import { uiEventRecorder } from '../../../amazonq/util/eventRecorder'
 import { globals } from '../../../shared'
 import { telemetry } from '../../../shared/telemetry'
+import { isSsoConnection } from '../../../auth/connection'
 
 export interface ChatControllerMessagePublishers {
     readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -656,12 +655,24 @@ export class ChatController {
                 request
             )}`
         )
-        let response: GenerateAssistantResponseCommandOutput | undefined = undefined
+        let response: MessengerResponseType | undefined = undefined
         session.createNewTokenSource()
         try {
             this.messenger.sendInitalStream(tabID, triggerID)
             this.telemetryHelper.setConversationStreamStartTime(tabID)
-            response = await session.chat(request)
+            if (isSsoConnection(AuthUtil.instance.conn)) {
+                const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)
+                response = {
+                    $metadata: $metadata,
+                    message: generateAssistantResponseResponse,
+                }
+            } else {
+                const { $metadata, sendMessageResponse } = await session.chatIam(request as SendMessageRequest)
+                response = {
+                    $metadata: $metadata,
+                    message: sendMessageResponse,
+                }
+            }
             this.telemetryHelper.recordEnterFocusConversation(triggerEvent.tabID)
             this.telemetryHelper.recordStartConversation(triggerEvent, triggerPayload)
 

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -13,10 +13,12 @@ import {
     QuickActionMessage,
 } from '../../../view/connector/connector'
 import { EditorContextCommandType } from '../../../commands/registerCommands'
+import { ChatResponseStream as qdevChatResponseStream } from '@amzn/amazon-q-developer-streaming-client'
 import {
+    ChatResponseStream as cwChatResponseStream,
     CodeWhispererStreamingServiceException,
-    GenerateAssistantResponseCommandOutput,
     SupplementaryWebLink,
+    __MetadataBearer,
 } from '@amzn/codewhisperer-streaming'
 import { ChatMessage, ErrorMessage, FollowUp, Suggestion } from '../../../view/connector/connector'
 import { ChatSession } from '../../../clients/chat/v0/chat'
@@ -37,6 +39,10 @@ import { extractAuthFollowUp } from '../../../../amazonq/util/authUtils'
 
 export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
 
+export type MessengerResponseType = __MetadataBearer & {
+    message?: AsyncIterable<cwChatResponseStream | qdevChatResponseStream>
+}
+
 export class Messenger {
     public constructor(
         private readonly dispatcher: AppToWebViewMessageDispatcher,
@@ -103,7 +109,7 @@ export class Messenger {
     }
 
     public async sendAIResponse(
-        response: GenerateAssistantResponseCommandOutput,
+        response: MessengerResponseType,
         session: ChatSession,
         tabID: string,
         triggerID: string,
@@ -116,7 +122,7 @@ export class Messenger {
         let relatedSuggestions: Suggestion[] = []
         let codeBlockLanguage: string = 'plaintext'
 
-        if (response.generateAssistantResponseResponse === undefined) {
+        if (response.message === undefined) {
             throw new ToolkitError(
                 `Empty response from CodeWhisperer Streaming service. Request ID: ${response.$metadata.requestId}`
             )
@@ -133,7 +139,7 @@ export class Messenger {
         const eventCounts = new Map<string, number>()
         waitUntil(
             async () => {
-                for await (const chatEvent of response.generateAssistantResponseResponse!) {
+                for await (const chatEvent of response.message!) {
                     for (const key of keys(chatEvent)) {
                         if ((chatEvent[key] as any) !== undefined) {
                             eventCounts.set(key, (eventCounts.get(key) ?? 0) + 1)

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -0,0 +1,25 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { QDeveloperStreaming } from '@amzn/amazon-q-developer-streaming-client'
+import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer'
+import { getUserAgent } from '../telemetry/util'
+import { ConfiguredRetryStrategy } from '@smithy/util-retry'
+import { AuthUtil } from '../../codewhisperer/util/authUtil'
+
+// Create a client for featureDev streaming based off of aws sdk v3
+export async function createQDeveloperStreamingClient(): Promise<QDeveloperStreaming> {
+    const cwsprConfig = getCodewhispererConfig()
+    const credentials = await AuthUtil.instance.getCredentials()
+    const streamingClient = new QDeveloperStreaming({
+        region: cwsprConfig.region,
+        endpoint: cwsprConfig.endpoint,
+        credentials: credentials,
+        customUserAgent: getUserAgent(),
+        // SETTING max attempts to 0 FOR BETA. RE-ENABLE FOR RE-INVENT
+        // Implement exponential back off starting with a base of 500ms (500 + attempt^10)
+        retryStrategy: new ConfiguredRetryStrategy(0, (attempt: number) => 500 + attempt ** 10),
+    })
+    return streamingClient
+}

packages/core/src/shared/clients/qDeveloperChatClient.ts

@@ -709,7 +709,7 @@
                 {
                     "id": "aws.amazonq.codewhisperer",
                     "name": "%AWS.amazonq.codewhisperer.title%",
-                    "when": "!isCloud9 && !aws.isSageMaker && !aws.toolkit.amazonq.dismissed && !aws.explorer.showAuthView"
+                    "when": "!isCloud9 && !aws.toolkit.amazonq.dismissed && !aws.explorer.showAuthView"
                 },
                 {
                     "id": "aws.explorer",