feat(server-utils): Record Bedrock prompts, tools, and responses

dev-packages/node-integration-tests/suites/aws-serverless/bedrock/instrument-gen-ai.mjs

@@ -0,0 +1,14 @@
+import * as Sentry from '@sentry/aws-serverless';
+import { loggingTransport } from '@sentry-internal/node-integration-tests';
+
+const genAI = process.env.GEN_AI_DATA_COLLECTION && JSON.parse(process.env.GEN_AI_DATA_COLLECTION);
+const awsOptions = process.env.AWS_INTEGRATION_OPTIONS && JSON.parse(process.env.AWS_INTEGRATION_OPTIONS);
+
+Sentry.init({
+  dsn: 'https://public@dsn.ingest.sentry.io/1337',
+  release: '1.0',
+  tracesSampleRate: 1.0,
+  dataCollection: genAI ? { genAI } : undefined,
+  integrations: awsOptions ? [Sentry.awsIntegration(awsOptions)] : undefined,
+  transport: loggingTransport,
+});

dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario-stream.mjs

@@ -0,0 +1,194 @@
+/* eslint-disable no-bitwise -- CRC32 for the AWS event stream framing */
+import * as Sentry from '@sentry/aws-serverless';
+import {
+  BedrockRuntimeClient,
+  ConverseStreamCommand,
+  InvokeModelWithResponseStreamCommand,
+} from '@aws-sdk/client-bedrock-runtime';
+// Force the HTTP/1 handler so `nock` can intercept the request instead of hitting real AWS.
+import { NodeHttpHandler } from '@smithy/node-http-handler';
+import nock from 'nock';
+import { Buffer } from 'node:buffer';
+
+nock.disableNetConnect();
+
+const region = 'us-east-1';
+const credentials = { accessKeyId: 'aws-test-key', secretAccessKey: 'aws-test-secret' };
+const host = `https://bedrock-runtime.${region}.amazonaws.com`;
+const modelId = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
+
+// Streaming responses use the AWS event stream framing: a prelude with the total and header
+// lengths, string headers, the payload, and CRC32 checksums the SDK verifies.
+// https://docs.aws.amazon.com/AmazonS3/latest/API/RESTSelectObjectAppendix.html
+const CRC32_TABLE = Array.from({ length: 256 }, (_, n) => {
+  let c = n;
+  for (let k = 0; k < 8; k++) {
+    c = c & 1 ? 0xedb88320 ^ (c >>> 1) : c >>> 1;
+  }
+  return c >>> 0;
+});
+
+function crc32(bytes) {
+  let crc = 0xffffffff;
+  for (const byte of bytes) {
+    crc = CRC32_TABLE[(crc ^ byte) & 0xff] ^ (crc >>> 8);
+  }
+  return (crc ^ 0xffffffff) >>> 0;
+}
+
+function uint32(value) {
+  const buffer = Buffer.alloc(4);
+  buffer.writeUInt32BE(value);
+  return buffer;
+}
+
+function encodeMessage(headers, payload) {
+  const encodedHeaders = Buffer.concat(
+    Object.entries(headers).map(([name, value]) => {
+      const nameBytes = Buffer.from(name);
+      const valueBytes = Buffer.from(value);
+      const valueLength = Buffer.alloc(2);
+      valueLength.writeUInt16BE(valueBytes.length);
+      // 7 is the header value type for strings.
+      return Buffer.concat([Buffer.from([nameBytes.length]), nameBytes, Buffer.from([7]), valueLength, valueBytes]);
+    }),
+  );
+  const body = Buffer.from(JSON.stringify(payload));
+  const prelude = Buffer.concat([uint32(12 + encodedHeaders.length + body.length + 4), uint32(encodedHeaders.length)]);
+  const message = Buffer.concat([prelude, uint32(crc32(prelude)), encodedHeaders, body]);
+  return Buffer.concat([message, uint32(crc32(message))]);
+}
+
+function event(eventType, payload) {
+  return encodeMessage(
+    { ':message-type': 'event', ':event-type': eventType, ':content-type': 'application/json' },
+    payload,
+  );
+}
+
+function exception(exceptionType, payload) {
+  return encodeMessage(
+    { ':message-type': 'exception', ':exception-type': exceptionType, ':content-type': 'application/json' },
+    payload,
+  );
+}
+
+function invokeModelChunk(chunk) {
+  return event('chunk', { bytes: Buffer.from(JSON.stringify(chunk)).toString('base64') });
+}
+
+function mockStream(path, messages) {
+  nock(host).post(path).reply(200, Buffer.concat(messages), { 'content-type': 'application/vnd.amazon.eventstream' });
+}
+
+function createClient() {
+  return new BedrockRuntimeClient({ region, credentials, maxAttempts: 1, requestHandler: new NodeHttpHandler() });
+}
+
+async function converseStream() {
+  mockStream(/\/model\/.*\/converse-stream$/, [
+    event('messageStart', { role: 'assistant' }),
+    event('contentBlockDelta', { contentBlockIndex: 0, delta: { text: 'Hello ' } }),
+    event('contentBlockDelta', { contentBlockIndex: 0, delta: { text: 'from Bedrock' } }),
+    event('contentBlockStop', { contentBlockIndex: 0 }),
+    event('contentBlockStart', {
+      contentBlockIndex: 1,
+      start: { toolUse: { toolUseId: 'tooluse_1', name: 'get_weather' } },
+    }),
+    event('contentBlockDelta', { contentBlockIndex: 1, delta: { toolUse: { input: '{"city":' } } }),
+    event('contentBlockDelta', { contentBlockIndex: 1, delta: { toolUse: { input: '"Vienna"}' } } }),
+    event('contentBlockStop', { contentBlockIndex: 1 }),
+    event('messageStop', { stopReason: 'tool_use' }),
+    event('metadata', { usage: { inputTokens: 12, outputTokens: 8, totalTokens: 20 }, metrics: { latencyMs: 100 } }),
+  ]);
+
+  const response = await createClient().send(
+    new ConverseStreamCommand({
+      modelId,
+      system: [{ text: 'You are a helpful assistant.' }],
+      messages: [{ role: 'user', content: [{ text: 'Hello' }] }],
+    }),
+  );
+
+  let text = '';
+  for await (const item of response.stream) {
+    text += item.contentBlockDelta?.delta?.text ?? '';
+  }
+  if (text !== 'Hello from Bedrock') {
+    throw new Error(`Unexpected streamed text: ${text}`);
+  }
+}
+
+async function invokeModelWithResponseStream() {
+  mockStream(/\/model\/.*\/invoke-with-response-stream$/, [
+    invokeModelChunk({
+      type: 'message_start',
+      message: { id: 'msg_bdrk_2', model: 'claude-3-5-sonnet-20240620', usage: { input_tokens: 15, output_tokens: 1 } },
+    }),
+    invokeModelChunk({ type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } }),
+    invokeModelChunk({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Streamed ' } }),
+    invokeModelChunk({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'reply' } }),
+    invokeModelChunk({ type: 'content_block_stop', index: 0 }),
+    invokeModelChunk({ type: 'message_delta', delta: { stop_reason: 'end_turn' }, usage: { output_tokens: 9 } }),
+    invokeModelChunk({ type: 'message_stop' }),
+  ]);
+
+  const response = await createClient().send(
+    new InvokeModelWithResponseStreamCommand({
+      modelId,
+      contentType: 'application/json',
+      body: JSON.stringify({
+        anthropic_version: 'bedrock-2023-05-31',
+        max_tokens: 100,
+        messages: [{ role: 'user', content: 'Stream please' }],
+      }),
+    }),
+  );
+
+  let chunks = 0;
+  for await (const item of response.body) {
+    if (item.chunk?.bytes) {
+      chunks++;
+    }
+  }
+  if (chunks !== 7) {
+    throw new Error(`Unexpected chunk count: ${chunks}`);
+  }
+}
+
+async function failingConverseStream() {
+  mockStream(/\/model\/.*\/converse-stream$/, [
+    event('messageStart', { role: 'assistant' }),
+    event('contentBlockDelta', { contentBlockIndex: 0, delta: { text: 'Partial' } }),
+    exception('modelStreamErrorException', { message: 'The model stream failed', originalStatusCode: 500 }),
+  ]);
+
+  const response = await createClient().send(
+    new ConverseStreamCommand({
+      modelId: 'amazon.nova-pro-v1:0',
+      messages: [{ role: 'user', content: [{ text: 'Fail please' }] }],
+    }),
+  );
+
+  try {
+    for await (const _ of response.stream) {
+      // drain
+    }
+  } catch (error) {
+    if (error.name !== 'ModelStreamErrorException') {
+      throw error;
+    }
+    return;
+  }
+  throw new Error('Expected the stream to fail');
+}
+
+async function run() {
+  await Sentry.startSpan({ name: 'Test Transaction' }, async () => {
+    await converseStream();
+    await invokeModelWithResponseStream();
+    await failingConverseStream();
+  });
+}
+
+run();

dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario.mjs

@@ -23,8 +23,16 @@ async function converse() {
     .reply(
       200,
       JSON.stringify({
-        output: { message: { role: 'assistant', content: [{ text: 'Hello from Bedrock' }] } },
-        stopReason: 'end_turn',
+        output: {
+          message: {
+            role: 'assistant',
+            content: [
+              { text: 'Hello from Bedrock' },
+              { toolUse: { toolUseId: 'tooluse_1', name: 'get_weather', input: { city: 'Vienna' } } },
+            ],
+          },
+        },
+        stopReason: 'tool_use',
         usage: { inputTokens: 12, outputTokens: 8, totalTokens: 20 },
       }),
       { 'content-type': 'application/json' },
@@ -33,8 +41,20 @@ async function converse() {
   await client.send(
     new ConverseCommand({
       modelId: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
+      system: [{ text: 'You are a helpful assistant.' }],
       messages: [{ role: 'user', content: [{ text: 'Hello' }] }],
       inferenceConfig: { maxTokens: 100, temperature: 0.5, topP: 0.9 },
+      toolConfig: {
+        tools: [
+          {
+            toolSpec: {
+              name: 'get_weather',
+              description: 'Get the weather for a city',
+              inputSchema: { json: { type: 'object', properties: { city: { type: 'string' } } } },
+            },
+          },
+        ],
+      },
     }),
   );
 }
@@ -52,6 +72,8 @@ async function invokeModel() {
     .reply(
       200,
       JSON.stringify({
+        id: 'msg_bdrk_1',
+        model: 'claude-3-5-sonnet-20240620',
         content: [{ type: 'text', text: 'Hello from Bedrock' }],
         stop_reason: 'end_turn',
         usage: { input_tokens: 15, output_tokens: 9 },
@@ -68,7 +90,29 @@ async function invokeModel() {
         max_tokens: 100,
         temperature: 0.5,
         top_p: 0.9,
-        messages: [{ role: 'user', content: 'Hello' }],
+        system: 'You are a helpful assistant.',
+        messages: [
+          {
+            role: 'user',
+            content: [
+              { type: 'text', text: 'Hello' },
+              { type: 'image', source: { type: 'base64', media_type: 'image/png', data: 'iVBORw0KGgo=' } },
+            ],
+          },
+          { role: 'assistant', content: [{ type: 'tool_use', id: 'toolu_1', name: 'get_report', input: {} }] },
+          {
+            role: 'user',
+            content: [
+              {
+                type: 'tool_result',
+                tool_use_id: 'toolu_1',
+                content: [
+                  { type: 'document', source: { type: 'base64', media_type: 'application/pdf', data: 'JVBERi0xLjQ=' } },
+                ],
+              },
+            ],
+          },
+        ],
       }),
     }),
   );