Support Bedrock Anthropic Stream for Messages API (#1226)

docs/pages/docs/api-reference/providers/aws-bedrock-stream.mdx

@@ -2,9 +2,10 @@
 
 The AWS Bedrock stream functions are utilties that transform the outputs from the [AWS Bedrock API](https://docs.aws.amazon.com/bedrock/latest/APIReference/welcome.html) into a `ReadableStream`. They use `AIStream` under the hood and handle parsing Bedrock's response.
 
-There are currently 3 AWS Bedrock stream wrappers:
+There are currently 4 AWS Bedrock stream wrappers:
 
 - `AWSBedrockAnthropicStream` for Anthropic models
+- `AWSBedrockAnthropicMessagesStream` for Anthropic models using Messages API
 - `AWSBedrockCohereStream` for Cohere models
 - `AWSBedrockLlama2Stream` for Llama 2 models
 
@@ -74,6 +75,59 @@ export async function POST(req: Request) {
 
 In this example, the `AWSBedrockAnthropicStream` function transforms the text generation stream from the Bedrock into a ReadableStream of parsed result. This allows clients to consume AI outputs in real-time as they're generated, instead of waiting for the complete response.
 
+## Example: Anthropic Models using Messages API
+
+The `AWSBedrockAnthropicMessagesStream` function can be coupled with a call to an Anthropic model using the [Messages API](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html). This stream can then facilitate the real-time consumption of AI outputs as they're being generated.
+
+Here's a step-by-step example of how to implement this in Next.js:
+
+```js filename="app/api/chat/route.ts"
+import {
+  BedrockRuntimeClient,
+  InvokeModelWithResponseStreamCommand,
+} from '@aws-sdk/client-bedrock-runtime';
+import { AWSBedrockAnthropicMessagesStream, StreamingTextResponse } from 'ai';
+import { experimental_buildAnthropicMessages } from 'ai/prompts';
+
+// IMPORTANT! Set the runtime to edge
+export const runtime = 'edge';
+
+const bedrockClient = new BedrockRuntimeClient({
+  region: process.env.AWS_REGION ?? 'us-east-1',
+  credentials: {
+    accessKeyId: process.env.AWS_ACCESS_KEY_ID ?? '',
+    secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY ?? '',
+  },
+});
+
+export async function POST(req: Request) {
+  // Extract the `prompt` from the body of the request
+  const { messages } = await req.json();
+
+  // Ask Claude for a streaming chat completion given the prompt
+  const bedrockResponse = await bedrockClient.send(
+    new InvokeModelWithResponseStreamCommand({
+      modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
+      contentType: 'application/json',
+      accept: 'application/json',
+      body: JSON.stringify({
+        anthropic_version: 'bedrock-2023-05-31',
+        message: experimental_buildAnthropicMessages(messages),
+        max_tokens: 300,
+      }),
+    }),
+  );
+
+  // Convert the response into a friendly text-stream
+  const stream = AWSBedrockAnthropicMessagesStream(bedrockResponse);
+
+  // Respond with the stream
+  return new StreamingTextResponse(stream);
+}
+```
+
+In this example, the `AWSBedrockAnthropicStream` function transforms the text generation stream from the Bedrock into a ReadableStream of parsed result. This allows clients to consume AI outputs in real-time as they're generated, instead of waiting for the complete response.
+
 ## Example: Cohere Models
 
 The `AWSBedrockCohereStream` function can be coupled with a call to a Cohere model using the Bedrock API to generate a readable stream of the completion. This stream can then facilitate the real-time consumption of AI outputs as they're being generated.

examples/next-aws-bedrock/app/api/chat-anthropic-v3/route.ts

@@ -0,0 +1,42 @@
+import {
+  BedrockRuntimeClient,
+  InvokeModelWithResponseStreamCommand,
+} from '@aws-sdk/client-bedrock-runtime';
+import { AWSBedrockAnthropicMessagesStream, StreamingTextResponse } from 'ai';
+import { experimental_buildAnthropicMessages } from 'ai/prompts';
+
+// IMPORTANT! Set the runtime to edge
+export const runtime = 'edge';
+
+const bedrockClient = new BedrockRuntimeClient({
+  region: process.env.AWS_REGION ?? 'us-east-1',
+  credentials: {
+    accessKeyId: process.env.AWS_ACCESS_KEY_ID ?? '',
+    secretAccessKey: process.env.AWS_SECRET_ACCESS_KEY ?? '',
+  },
+});
+
+export async function POST(req: Request) {
+  // Extract the `prompt` from the body of the request
+  const { messages } = await req.json();
+
+  // Ask Claude for a streaming chat completion given the prompt
+  const bedrockResponse = await bedrockClient.send(
+    new InvokeModelWithResponseStreamCommand({
+      modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
+      contentType: 'application/json',
+      accept: 'application/json',
+      body: JSON.stringify({
+        anthropic_version: 'bedrock-2023-05-31',
+        message: experimental_buildAnthropicMessages(messages),
+        max_tokens: 300,
+      }),
+    }),
+  );
+
+  // Convert the response into a friendly text-stream
+  const stream = AWSBedrockAnthropicMessagesStream(bedrockResponse);
+
+  // Respond with the stream
+  return new StreamingTextResponse(stream);
+}

packages/core/prompts/anthropic.ts

@@ -18,3 +18,22 @@ export function experimental_buildAnthropicPrompt(
     }) + '\n\nAssistant:'
   );
 }
+
+/**
+ * A prompt constructor for Anthropic V3 models which require Messages API.
+ * Does not support message with image content
+ * @see https://docs.anthropic.com/claude/reference/messages_post
+ */
+export function experimental_buildAnthropicMessages(
+  messages: Pick<Message, 'content' | 'role'>[],
+) {
+  return messages.map(({ content, role }) => {
+    if (!['assistant', 'user'].includes(role)) {
+      throw new Error(`Cannot use ${role} on Anthropic V3 Messages API`);
+    }
+    return {
+      role,
+      content: [{ type: 'text', text: content }],
+    };
+  });
+}

packages/core/streams/aws-bedrock-stream.test.ts

@@ -1,11 +1,13 @@
 import { StreamingTextResponse, experimental_StreamData } from '.';
 import {
   bedrockAnthropicChunks,
+  bedrockAnthropicV3Chunks,
   bedrockCohereChunks,
   bedrockLlama2Chunks,
 } from '../tests/snapshots/aws-bedrock';
 import { readAllChunks } from '../tests/utils/mock-client';
 import {
+  AWSBedrockAnthropicMessagesStream,
   AWSBedrockAnthropicStream,
   AWSBedrockCohereStream,
   AWSBedrockLlama2Stream,
@@ -98,6 +100,21 @@ describe('AWS Bedrock', () => {
     });
   });
 
+  describe('AnthropicV3', () => {
+    it('should be able to parse SSE and receive the streamed response', async () => {
+      const bedrockResponse = simulateBedrockResponse(bedrockAnthropicV3Chunks);
+      const stream = AWSBedrockAnthropicMessagesStream(bedrockResponse);
+      const response = new StreamingTextResponse(stream);
+
+      expect(await readAllChunks(response)).toEqual([
+        ' Hello',
+        ',',
+        ' world',
+        '.',
+      ]);
+    });
+  });
+
   describe('Cohere', () => {
     it('should be able to parse SSE and receive the streamed response', async () => {
       const bedrockResponse = simulateBedrockResponse(bedrockCohereChunks);

packages/core/streams/aws-bedrock-stream.ts

@@ -31,6 +31,13 @@ async function* asDeltaIterable(
   }
 }
 
+export function AWSBedrockAnthropicMessagesStream(
+  response: AWSBedrockResponse,
+  callbacks?: AIStreamCallbacksAndOptions,
+): ReadableStream {
+  return AWSBedrockStream(response, callbacks, chunk => chunk.delta?.text);
+}
+
 export function AWSBedrockAnthropicStream(
   response: AWSBedrockResponse,
   callbacks?: AIStreamCallbacksAndOptions,

packages/core/tests/snapshots/aws-bedrock.ts

@@ -5,6 +5,29 @@ export const bedrockAnthropicChunks = [
   { completion: '.', stop_reason: 'stop_sequence', stop: '\n\nHuman:' },
 ];
 
+export const bedrockAnthropicV3Chunks = [
+  {
+    type: 'content_block_delta',
+    index: 0,
+    delta: { type: 'text_delta', text: ' Hello' },
+  },
+  {
+    type: 'content_block_delta',
+    index: 0,
+    delta: { type: 'text_delta', text: ',' },
+  },
+  {
+    type: 'content_block_delta',
+    index: 0,
+    delta: { type: 'text_delta', text: ' world' },
+  },
+  {
+    type: 'content_block_delta',
+    index: 0,
+    delta: { type: 'text_delta', text: '.' },
+  },
+];
+
 export const bedrockCohereChunks = [
   {
     generations: [