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]>
aws-amplify · Oct 4, 2024 · 0a5e51c · 0a5e51c
1 parent 300a72d
commit 0a5e51c
Show file tree

Hide file tree

Showing 22 changed files with 375 additions and 182 deletions.
diff --git a/.changeset/eleven-snails-move.md b/.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
diff --git a/.prettierignore b/.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
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/ai-constructs/API.md b/packages/ai-constructs/API.md
@@ -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)

diff --git a/packages/ai-constructs/package.json b/packages/ai-constructs/package.json
@@ -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": {

diff --git a/packages/ai-constructs/src/conversation/conversation_handler_construct.test.ts b/packages/ai-constructs/src/conversation/conversation_handler_construct.test.ts
@@ -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);

diff --git a/packages/ai-constructs/src/conversation/conversation_handler_construct.ts b/packages/ai-constructs/src/conversation/conversation_handler_construct.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,
+ },
+ });
+ };
 }
diff --git a/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.ts b/packages/ai-constructs/src/conversation/runtime/bedrock_converse_adapter.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: [

diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_turn_executor.test.ts b/packages/ai-constructs/src/conversation/runtime/conversation_turn_executor.test.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(

diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_turn_executor.ts b/packages/ai-constructs/src/conversation/runtime/conversation_turn_executor.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();
 

diff --git a/packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.ts b/packages/ai-constructs/src/conversation/runtime/conversation_turn_response_sender.ts
@@ -23,18 +23,17 @@ export class ConversationTurnResponseSender {
  private readonly graphqlRequestExecutor = new GraphqlRequestExecutor(
  event.graphqlApiEndpoint,
  event.request.headers.authorization
- )
+ ),
+ private readonly logger = console
  ) {}
 
  sendResponse = async (message: ContentBlock[]) => {
- const { query, variables } = this.createMutationRequest(message);
+ const responseMutationRequest = this.createMutationRequest(message);
+ this.logger.debug('Sending response mutation:', responseMutationRequest);
  await this.graphqlRequestExecutor.executeGraphql<
  MutationResponseInput,
  void
- >({
- query,
- variables,
- });
+ >(responseMutationRequest);
  };
 
  private createMutationRequest = (content: ContentBlock[]) => {

diff --git a/packages/ai-constructs/tsconfig.json b/packages/ai-constructs/tsconfig.json
@@ -9,5 +9,10 @@
  "outDir": "lib",
  "allowJs": true
  },
- "references": [{ "path": "../plugin-types" }]
+ "references": [
+ { "path": "../backend-output-schemas" },
+ { "path": "../platform-core" },
+ { "path": "../plugin-types" },
+ { "path": "../backend-output-storage" }
+ ]
 }
diff --git a/packages/backend-ai/src/conversation/factory.test.ts b/packages/backend-ai/src/conversation/factory.test.ts
@@ -158,8 +158,8 @@ void describe('ConversationHandlerFactory', () => {
  });
  factory.getInstance(getInstanceProps);
  const template = Template.fromStack(rootStack);
- const outputValue =
- template.findOutputs('definedFunctions').definedFunctions.Value;
+ const outputValue = template.findOutputs('definedConversationHandlers')
+ .definedConversationHandlers.Value;
  assert.deepStrictEqual(outputValue, {
  ['Fn::Join']: [
  '',