From faf84848e42ec02994003283bf6178aa55233531 Mon Sep 17 00:00:00 2001
From: Robert Craigie <robert@craigie.dev>
Date: Tue, 22 Oct 2024 11:49:12 +0100
Subject: [PATCH] feat(bedrock): add beta.messages.create() method

---
 packages/bedrock-sdk/src/client.ts | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/packages/bedrock-sdk/src/client.ts b/packages/bedrock-sdk/src/client.ts
index b275ce7d..c2ab6ebc 100644
--- a/packages/bedrock-sdk/src/client.ts
+++ b/packages/bedrock-sdk/src/client.ts
@@ -76,6 +76,7 @@ export class AnthropicBedrock extends Core.APIClient {
 
   messages: Resources.Messages = new Resources.Messages(this);
   completions: Resources.Completions = new Resources.Completions(this);
+  beta: BetaResource = makeBetaResource(this);
 
   protected override defaultQuery(): Core.DefaultQuery | undefined {
     return this._options.defaultQuery;
@@ -120,6 +121,13 @@ export class AnthropicBedrock extends Core.APIClient {
       if (!options.body['anthropic_version']) {
         options.body['anthropic_version'] = DEFAULT_VERSION;
       }
+
+      if (options.headers && !options.body['anthropic_beta']) {
+        const betas = Core.getHeader(options.headers, 'anthropic-beta');
+        if (betas != null) {
+          options.body['anthropic_beta'] = betas.split(',');
+        }
+      }
     }
 
     if (MODEL_ENDPOINTS.has(options.path) && options.method === 'post') {
@@ -143,3 +151,22 @@ export class AnthropicBedrock extends Core.APIClient {
     return super.buildRequest(options);
   }
 }
+
+/**
+ * The Bedrock API does not currently support prompt caching or the Batch API.
+ */
+type BetaResource = Omit<Resources.Beta, 'promptCaching' | 'messages'> & {
+  messages: Omit<Resources.Beta['messages'], 'batches'>;
+};
+
+function makeBetaResource(client: AnthropicBedrock): BetaResource {
+  const resource = new Resources.Beta(client);
+
+  // @ts-expect-error we're deleting non-optional properties
+  delete resource.promptCaching;
+
+  // @ts-expect-error we're deleting non-optional properties
+  delete resource.messages.batches;
+
+  return resource;
+}