feat(bedrock): add messages API (#305)

anthropics · Mar 4, 2024 · 8b7f89e · 8b7f89e
1 parent 459956a
commit 8b7f89e
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 30 deletions.
diff --git a/packages/bedrock-sdk/README.md b/packages/bedrock-sdk/README.md
@@ -18,7 +18,6 @@ yarn add @anthropic-ai/bedrock-sdk
 
 <!-- prettier-ignore -->
 ```js
-import Anthropic from '@anthropic-ai/sdk';
 import { AnthropicBedrock } from '@anthropic-ai/bedrock-sdk';
 
 // Note: this assumes you have configured AWS credentials in a way
@@ -29,16 +28,17 @@ import { AnthropicBedrock } from '@anthropic-ai/bedrock-sdk';
 const anthropic = new AnthropicBedrock();
 
 async function main() {
-  const completion = await anthropic.completions.create({
-    model: 'anthropic.claude-3-opus-20240229-v1:0',
-    prompt: `${Anthropic.HUMAN_PROMPT} how does a court case get to the Supreme Court? ${Anthropic.AI_PROMPT}`,
-    stop_sequences: [Anthropic.HUMAN_PROMPT],
-    max_tokens_to_sample: 800,
-    temperature: 0.5,
-    top_k: 250,
-    top_p: 0.5,
+  const message = await anthropic.messages.create({
+    model: 'anthropic.claude-3-sonnet-20240229-v1:0',
+    messages: [
+      {
+        role: 'user',
+        content: 'Hello!',
+      },
+    ],
+    max_tokens: 1024,
   });
-  console.log(completion);
+  console.log(message);
 }
 
 main();

diff --git a/packages/bedrock-sdk/examples/demo.ts b/packages/bedrock-sdk/examples/demo.ts
@@ -1,6 +1,5 @@
 #!/usr/bin/env -S npm run tsn -T
 
-import Anthropic from '@anthropic-ai/sdk';
 import { AnthropicBedrock } from '@anthropic-ai/bedrock-sdk';
 
 // Note: this assumes you have configured AWS credentials in a way
@@ -11,16 +10,17 @@ import { AnthropicBedrock } from '@anthropic-ai/bedrock-sdk';
 const anthropic = new AnthropicBedrock();
 
 async function main() {
-  const completion = await anthropic.completions.create({
-    model: 'anthropic.claude-3-opus-20240229-v1:0',
-    prompt: `${Anthropic.HUMAN_PROMPT} how does a court case get to the Supreme Court? ${Anthropic.AI_PROMPT}`,
-    stop_sequences: [Anthropic.HUMAN_PROMPT],
-    max_tokens_to_sample: 800,
-    temperature: 0.5,
-    top_k: 250,
-    top_p: 0.5,
+  const message = await anthropic.messages.create({
+    model: 'anthropic.claude-3-sonnet-20240229-v1:0',
+    messages: [
+      {
+        role: 'user',
+        content: 'Hello!',
+      },
+    ],
+    max_tokens: 1024,
   });
-  console.log(completion);
+  console.log(message);
 }
 
 main();
diff --git a/packages/bedrock-sdk/examples/streaming.ts b/packages/bedrock-sdk/examples/streaming.ts
@@ -1,6 +1,5 @@
 #!/usr/bin/env -S npm run tsn -T
 
-import Anthropic from '@anthropic-ai/sdk';
 import { AnthropicBedrock } from '@anthropic-ai/bedrock-sdk';
 
 // Note: this assumes you have configured AWS credentials in a way
@@ -11,18 +10,24 @@ import { AnthropicBedrock } from '@anthropic-ai/bedrock-sdk';
 const client = new AnthropicBedrock();
 
 async function main() {
-  const question = 'Hey Claude! How can I recursively list all files in a directory in Rust?';
-
-  const stream = await client.completions.create({
-    prompt: `${Anthropic.HUMAN_PROMPT}${question}${Anthropic.AI_PROMPT}:`,
-    model: 'anthropic.claude-v2:1',
+  const stream = await client.messages.create({
+    model: 'anthropic.claude-3-sonnet-20240229-v1:0',
+    messages: [
+      {
+        role: 'user',
+        content: 'Hello!',
+      },
+    ],
+    max_tokens: 1024,
     stream: true,
-    max_tokens_to_sample: 500,
   });
 
-  for await (const completion of stream) {
-    process.stdout.write(completion.completion);
+  for await (const event of stream) {
+    if (event.type === 'content_block_delta' && event.delta.type === 'text_delta') {
+      process.stdout.write(event.delta.text);
+    }
   }
+  process.stdout.write('\n');
 }
 
 main();
diff --git a/packages/bedrock-sdk/src/client.ts b/packages/bedrock-sdk/src/client.ts
@@ -5,6 +5,7 @@ import { getAuthHeaders } from './auth';
 import { Stream } from './streaming';
 
 const DEFAULT_VERSION = 'bedrock-2023-05-31';
+const MODEL_ENDPOINTS = new Set<string>(['/v1/complete', '/v1/messages']);
 
 export type ClientOptions = Omit<API.ClientOptions, 'apiKey' | 'authToken'> & {
   awsSecretKey?: string | null | undefined;
@@ -73,6 +74,7 @@ export class AnthropicBedrock extends Core.APIClient {
     this.awsSessionToken = awsSessionToken;
   }
 
+  messages: Resources.Messages = new Resources.Messages(this);
   completions: Resources.Completions = new Resources.Completions(this);
 
   protected override defaultQuery(): Core.DefaultQuery | undefined {
@@ -120,7 +122,7 @@ export class AnthropicBedrock extends Core.APIClient {
       }
     }
 
-    if (options.path === '/v1/complete' && options.method === 'post') {
+    if (MODEL_ENDPOINTS.has(options.path) && options.method === 'post') {
       if (!Core.isObj(options.body)) {
         throw new Error('Expected request body to be an object for post /v1/messages');
       }