feat: ✨ AWS Bedrock Provider tool streaming

core/llm/llms/Bedrock.ts

@@ -1,13 +1,19 @@
 import {
   BedrockRuntimeClient,
   ContentBlock,
+  ContentBlockDelta,
+  ContentBlockStart,
+  ContentBlockStartEvent,
   ConversationRole,
   ConverseStreamCommand,
   ConverseStreamCommandOutput,
   ImageFormat,
   InvokeModelCommand,
   Message,
+  ReasoningContentBlockDelta,
   ToolConfiguration,
+  ToolUseBlock,
+  ToolUseBlockDelta,
 } from "@aws-sdk/client-bedrock-runtime";
 import { fromNodeProviderChain } from "@aws-sdk/credential-providers";
 
@@ -25,15 +31,6 @@ interface ModelConfig {
   extractEmbeddings: (responseBody: any) => number[][];
 }
 
-/**
- * Interface for tool use state tracking
- */
-interface ToolUseState {
-  toolUseId: string;
-  name: string;
-  input: string;
-}
-
 /**
  * Interface for prompt caching metrics
  */
@@ -50,7 +47,6 @@ class Bedrock extends BaseLLM {
     profile: "bedrock",
   };
 
-  private _currentToolResponse: Partial<ToolUseState> | null = null;
   private _promptCachingMetrics: PromptCachingMetrics = {
     cacheReadInputTokens: 0,
     cacheWriteInputTokens: 0,
@@ -149,108 +145,99 @@ class Bedrock extends BaseLLM {
           console.log(`${JSON.stringify(chunk.metadata.usage)}`);
         }
 
-        if (chunk.contentBlockDelta?.delta) {
-          const delta: any = chunk.contentBlockDelta.delta;
-
+        const contentBlockDelta: ContentBlockDelta | undefined =
+          chunk.contentBlockDelta?.delta;
+        if (contentBlockDelta) {
           // Handle text content
-          if (chunk.contentBlockDelta.delta.text) {
+          if (contentBlockDelta.text) {
             yield {
               role: "assistant",
-              content: chunk.contentBlockDelta.delta.text,
-            };
-            continue;
-          }
-
-          // Handle text content
-          if ((chunk.contentBlockDelta.delta as any).reasoningContent?.text) {
-            yield {
-              role: "thinking",
-              content: (chunk.contentBlockDelta.delta as any).reasoningContent
-                .text,
+              content: contentBlockDelta.text,
             };
             continue;
           }
-
-          // Handle signature for thinking
-          if (delta.reasoningContent?.signature) {
+          if (contentBlockDelta.reasoningContent?.text) {
             yield {
               role: "thinking",
-              content: "",
-              signature: delta.reasoningContent.signature,
+              content: contentBlockDelta.reasoningContent.text,
             };
             continue;
           }
-
-          // Handle redacted thinking
-          if (delta.redactedReasoning?.data) {
+          if (contentBlockDelta.reasoningContent?.signature) {
             yield {
               role: "thinking",
               content: "",
-              redactedThinking: delta.redactedReasoning.data,
+              signature: contentBlockDelta.reasoningContent.signature,
             };
             continue;
           }
-
-          if (
-            chunk.contentBlockDelta.delta.toolUse?.input &&
-            this._currentToolResponse
-          ) {
-            // Append the new input to the existing string
-            if (this._currentToolResponse.input === undefined) {
-              this._currentToolResponse.input = "";
-            }
-            this._currentToolResponse.input +=
-              chunk.contentBlockDelta.delta.toolUse.input;
-            continue;
-          }
         }
 
-        if (chunk.contentBlockStart?.start) {
-          const start: any = chunk.contentBlockStart.start;
-          if (start.redactedReasoning) {
+        const reasoningDelta: ReasoningContentBlockDelta | undefined = chunk
+          .contentBlockDelta?.delta as ReasoningContentBlockDelta;
+        if (reasoningDelta) {
+          if (reasoningDelta.redactedContent) {
             yield {
               role: "thinking",
               content: "",
-              redactedThinking: start.redactedReasoning.data,
+              redactedThinking: reasoningDelta.text,
             };
             continue;
           }
+        }
 
-          const toolUse = chunk.contentBlockStart.start.toolUse;
-          if (toolUse?.toolUseId && toolUse?.name) {
-            this._currentToolResponse = {
-              toolUseId: toolUse.toolUseId,
-              name: toolUse.name,
-              input: "",
-            };
-          }
+        const toolUseBlockDelta: ToolUseBlockDelta | undefined = chunk
+          .contentBlockDelta?.delta?.toolUse as ToolUseBlockDelta;
+        const toolUseBlock: ToolUseBlock | undefined = chunk.contentBlockDelta
+          ?.delta?.toolUse as ToolUseBlock;
+        if (toolUseBlockDelta && toolUseBlock) {
+          yield {
+            role: "assistant",
+            content: "",
+            toolCalls: [
+              {
+                id: toolUseBlock.toolUseId,
+                type: "function",
+                function: {
+                  name: toolUseBlock.name,
+                  arguments: toolUseBlockDelta.input,
+                },
+              },
+            ],
+          };
           continue;
         }
 
-        if (chunk.contentBlockStop) {
-          if (this._currentToolResponse) {
-            yield {
-              role: "assistant",
-              content: "",
-              toolCalls: [
-                {
-                  id: this._currentToolResponse.toolUseId,
-                  type: "function",
-                  function: {
-                    name: this._currentToolResponse.name,
-                    arguments: this._currentToolResponse.input,
+        const contentBlockStart: ContentBlockStartEvent | undefined =
+          chunk.contentBlockStart as ContentBlockStartEvent;
+        if (contentBlockStart) {
+          const start: ContentBlockStart | undefined = chunk.contentBlockStart
+            ?.start as ContentBlockStart;
+          if (start) {
+            const toolUseBlock: ToolUseBlock | undefined =
+              start.toolUse as ToolUseBlock;
+            if (toolUseBlock?.toolUseId && toolUseBlock?.name) {
+              yield {
+                role: "assistant",
+                content: "",
+                toolCalls: [
+                  {
+                    id: toolUseBlock.toolUseId,
+                    type: "function",
+                    function: {
+                      name: toolUseBlock.name,
+                      arguments: "",
+                    },
                   },
-                },
-              ],
-            };
-            this._currentToolResponse = null;
+                ],
+              };
+              continue;
+            }
           }
-          continue;
         }
       }
     } catch (error: unknown) {
       // Clean up state and let the original error bubble up to the retry decorator
-      this._currentToolResponse = null;
       throw error;
     }
   }
@@ -348,6 +335,9 @@ class Bedrock extends BaseLLM {
               budget_tokens: options.reasoningBudgetTokens,
             }
           : undefined,
+        anthropic_beta: options.model.includes("claude")
+          ? ["fine-grained-tool-streaming-2025-05-14"]
+          : undefined,
       },
     };
   }