Stream conversation logs in sandbox (#2073)

* Stream conversation logs in sandbox

* fix that

* refactor that

* and that

* this works

* Undo the hack

* more tests

* that

* Update packages/sandbox/src/lambda_function_log_streamer.ts

Co-authored-by: Amplifiyer <[email protected]>

* use different output

* Revert "Undo the hack"

This reverts commit 3e8b590.

* use different outputs

* use different outputs

* Revert "Revert "Undo the hack""

This reverts commit 3c2d417.

* fix that

---------

Co-authored-by: Amplifiyer <[email protected]>

Author: Kamil Sobol

.changeset/eleven-snails-move.md

@@ -0,0 +1,9 @@
+---
+'@aws-amplify/ai-constructs': minor
+'@aws-amplify/backend-ai': minor
+'@aws-amplify/backend-output-schemas': minor
+'@aws-amplify/backend': patch
+'@aws-amplify/sandbox': patch
+---
+
+Stream conversation logs in sandbox

.prettierignore

@@ -1,4 +1,5 @@
 # Ignore artifacts:
+.amplify
 build
 coverage
 bin
@@ -10,6 +11,7 @@ verdaccio-cache
 expected-cdk-out
 .changeset/pre.json
 concurrent_workspace_script_cache.json
+packages/integration-tests/src/e2e-tests
 scripts/components/api-changes-validator/test-resources/working-directory
 /test-projects
-testDir
+testDir

package-lock.json

@@ -6,6 +6,8 @@
 
 /// <reference types="node" />
 
+import { AIConversationOutput } from '@aws-amplify/backend-output-schemas';
+import { BackendOutputStorageStrategy } from '@aws-amplify/plugin-types';
 import * as bedrock from '@aws-sdk/client-bedrock-runtime';
 import { Construct } from 'constructs';
 import { FunctionResources } from '@aws-amplify/plugin-types';
@@ -51,6 +53,7 @@ type ConversationHandlerFunctionProps = {
         modelId: string;
         region?: string;
     }>;
+    outputStorageStrategy?: BackendOutputStorageStrategy<AIConversationOutput>;
 };
 
 // @public (undocumented)

packages/ai-constructs/API.md

@@ -26,12 +26,15 @@
   },
   "license": "Apache-2.0",
   "dependencies": {
+    "@aws-amplify/backend-output-schemas": "^1.2.1",
+    "@aws-amplify/platform-core": "^1.1.0",
     "@aws-amplify/plugin-types": "^1.0.1",
     "@aws-sdk/client-bedrock-runtime": "^3.622.0",
     "@smithy/types": "^3.3.0",
     "json-schema-to-ts": "^3.1.1"
   },
   "devDependencies": {
+    "@aws-amplify/backend-output-storage": "^1.1.2",
     "typescript": "^5.0.0"
   },
   "peerDependencies": {

packages/ai-constructs/package.json

@@ -4,6 +4,7 @@ import { App, Stack } from 'aws-cdk-lib';
 import { ConversationHandlerFunction } from './conversation_handler_construct';
 import { Template } from 'aws-cdk-lib/assertions';
 import path from 'path';
+import { StackMetadataBackendOutputStorageStrategy } from '@aws-amplify/backend-output-storage';
 
 void describe('Conversation Handler Function construct', () => {
   void it('creates handler with log group with JWT token redacting policy', () => {
@@ -140,6 +141,56 @@ void describe('Conversation Handler Function construct', () => {
     });
   });
 
+  void it('does not store output if output strategy is absent', () => {
+    const app = new App();
+    const stack = new Stack(app);
+    new ConversationHandlerFunction(stack, 'conversationHandler', {
+      models: [
+        {
+          modelId: 'testModelId',
+        },
+      ],
+      outputStorageStrategy: undefined,
+    });
+    const template = Template.fromStack(stack);
+    const output = template.findOutputs(
+      'definedConversationHandlers'
+    ).definedConversationHandlers;
+    assert.ok(!output);
+  });
+
+  void it('stores output if output strategy is present', () => {
+    const app = new App();
+    const stack = new Stack(app);
+    new ConversationHandlerFunction(stack, 'conversationHandler', {
+      models: [
+        {
+          modelId: 'testModelId',
+        },
+      ],
+      outputStorageStrategy: new StackMetadataBackendOutputStorageStrategy(
+        stack
+      ),
+    });
+    const template = Template.fromStack(stack);
+    const outputValue = template.findOutputs('definedConversationHandlers')
+      .definedConversationHandlers.Value;
+    assert.deepStrictEqual(outputValue, {
+      'Fn::Join': [
+        '',
+        [
+          '["',
+          {
+            /* eslint-disable spellcheck/spell-checker */
+            Ref: 'conversationHandlerconversationHandlerFunction45BC2E1F',
+            /* eslint-enable spellcheck/spell-checker */
+          },
+          '"]',
+        ],
+      ],
+    });
+  });
+
   void it('throws if entry is not absolute', () => {
     const app = new App();
     const stack = new Stack(app);

packages/ai-constructs/src/conversation/conversation_handler_construct.test.ts

@@ -1,7 +1,15 @@
-import { FunctionResources, ResourceProvider } from '@aws-amplify/plugin-types';
-import { Duration, Stack } from 'aws-cdk-lib';
+import {
+  BackendOutputStorageStrategy,
+  FunctionResources,
+  ResourceProvider,
+} from '@aws-amplify/plugin-types';
+import { Duration, Stack, Tags } from 'aws-cdk-lib';
 import { Effect, PolicyStatement } from 'aws-cdk-lib/aws-iam';
-import { CfnFunction, Runtime as LambdaRuntime } from 'aws-cdk-lib/aws-lambda';
+import {
+  CfnFunction,
+  Runtime as LambdaRuntime,
+  LoggingFormat,
+} from 'aws-cdk-lib/aws-lambda';
 import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
 import {
   CustomDataIdentifier,
@@ -11,6 +19,11 @@ import {
 } from 'aws-cdk-lib/aws-logs';
 import { Construct } from 'constructs';
 import path from 'path';
+import { TagName } from '@aws-amplify/platform-core';
+import {
+  AIConversationOutput,
+  aiConversationOutputKey,
+} from '@aws-amplify/backend-output-schemas';
 
 const resourcesRoot = path.normalize(path.join(__dirname, 'runtime'));
 const defaultHandlerFilePath = path.join(resourcesRoot, 'default_handler.js');
@@ -21,6 +34,10 @@ export type ConversationHandlerFunctionProps = {
     modelId: string;
     region?: string;
   }>;
+  /**
+   * @internal
+   */
+  outputStorageStrategy?: BackendOutputStorageStrategy<AIConversationOutput>;
 };
 
 /**
@@ -53,6 +70,8 @@ export class ConversationHandlerFunction
       throw new Error('Entry must be absolute path');
     }
 
+    Tags.of(this).add(TagName.FRIENDLY_NAME, id);
+
     const conversationHandler = new NodejsFunction(
       this,
       `conversationHandlerFunction`,
@@ -67,6 +86,7 @@ export class ConversationHandlerFunction
           // For custom entry we do bundle SDK as we can't control version customer is coding against.
           bundleAwsSDK: !!this.props.entry,
         },
+        loggingFormat: LoggingFormat.JSON,
         logGroup: new LogGroup(this, 'conversationHandlerFunctionLogGroup', {
           retention: RetentionDays.INFINITE,
           dataProtectionPolicy: new DataProtectionPolicy({
@@ -105,5 +125,23 @@ export class ConversationHandlerFunction
         ) as CfnFunction,
       },
     };
+
+    this.storeOutput(this.props.outputStorageStrategy);
   }
+
+  /**
+   * Append conversation handler to defined functions.
+   */
+  private storeOutput = (
+    outputStorageStrategy:
+      | BackendOutputStorageStrategy<AIConversationOutput>
+      | undefined
+  ): void => {
+    outputStorageStrategy?.appendToBackendOutputList(aiConversationOutputKey, {
+      version: '1',
+      payload: {
+        definedConversationHandlers: this.resources.lambda.functionName,
+      },
+    });
+  };
 }

packages/ai-constructs/src/conversation/conversation_handler_construct.ts

@@ -41,7 +41,8 @@ export class BedrockConverseAdapter {
     eventToolsProvider = new ConversationTurnEventToolsProvider(event),
     private readonly messageHistoryRetriever = new ConversationMessageHistoryRetriever(
       event
-    )
+    ),
+    private readonly logger = console
   ) {
     this.executableTools = [
       ...eventToolsProvider.getEventTools(),
@@ -91,9 +92,16 @@ export class BedrockConverseAdapter {
         inferenceConfig: inferenceConfiguration,
         toolConfig,
       };
+      this.logger.info('Sending Bedrock Converse request');
+      this.logger.debug('Bedrock Converse request:', converseCommandInput);
       bedrockResponse = await this.bedrockClient.send(
         new ConverseCommand(converseCommandInput)
       );
+      this.logger.info(
+        `Received Bedrock Converse response, requestId=${bedrockResponse.$metadata.requestId}`,
+        bedrockResponse.usage
+      );
+      this.logger.debug('Bedrock Converse response:', bedrockResponse);
       if (bedrockResponse.output?.message) {
         messages.push(bedrockResponse.output?.message);
       }
@@ -194,7 +202,11 @@ export class BedrockConverseAdapter {
       );
     }
     try {
+      this.logger.info(`Invoking tool ${tool.name}`);
+      this.logger.debug('Tool input:', toolUseBlock.toolUse.input);
       const toolResponse = await tool.execute(toolUseBlock.toolUse.input);
+      this.logger.info(`Received response from ${tool.name} tool`);
+      this.logger.debug(toolResponse);
       return {
         role: 'user',
         content: [

packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts

@@ -47,9 +47,11 @@ void describe('Conversation turn executor', () => {
 
     const consoleErrorMock = mock.fn();
     const consoleLogMock = mock.fn();
+    const consoleDebugMock = mock.fn();
     const consoleMock = {
       error: consoleErrorMock,
       log: consoleLogMock,
+      debug: consoleDebugMock,
     } as unknown as Console;
 
     await new ConversationTurnExecutor(
@@ -100,9 +102,11 @@ void describe('Conversation turn executor', () => {
 
     const consoleErrorMock = mock.fn();
     const consoleLogMock = mock.fn();
+    const consoleDebugMock = mock.fn();
     const consoleMock = {
       error: consoleErrorMock,
       log: consoleLogMock,
+      debug: consoleDebugMock,
     } as unknown as Console;
 
     await assert.rejects(
@@ -164,9 +168,11 @@ void describe('Conversation turn executor', () => {
 
     const consoleErrorMock = mock.fn();
     const consoleLogMock = mock.fn();
+    const consoleDebugMock = mock.fn();
     const consoleMock = {
       error: consoleErrorMock,
       log: consoleLogMock,
+      debug: consoleDebugMock,
     } as unknown as Console;
 
     await assert.rejects(

packages/ai-constructs/src/conversation/runtime/conversation_turn_executor.test.ts

@@ -29,6 +29,7 @@ export class ConversationTurnExecutor {
       this.logger.log(
         `Handling conversation turn event, currentMessageId=${this.event.currentMessageId}, conversationId=${this.event.conversationId}`
       );
+      this.logger.debug('Event received:', this.event);
 
       const assistantResponse = await this.bedrockConverseAdapter.askBedrock();