From 660edbd9421c6f0f4cb891de4bf377459ad35ba0 Mon Sep 17 00:00:00 2001
From: Brace Sproul <braceasproul@gmail.com>
Date: Sun, 27 Oct 2024 13:58:30 -0700
Subject: [PATCH] fix(anthropic): Fixed streaming tool calls with Anthropic
 (#7081)

---
 libs/langchain-anthropic/src/chat_models.ts   |  6 +-
 .../src/utils/message_outputs.ts              | 16 +++++
 libs/langchain-anthropic/src/utils/tools.ts   | 58 -------------------
 3 files changed, 18 insertions(+), 62 deletions(-)

diff --git a/libs/langchain-anthropic/src/chat_models.ts b/libs/langchain-anthropic/src/chat_models.ts
index 3627aed6272c..f474324b37cd 100644
--- a/libs/langchain-anthropic/src/chat_models.ts
+++ b/libs/langchain-anthropic/src/chat_models.ts
@@ -32,7 +32,7 @@ import type {
 
 import { isLangChainTool } from "@langchain/core/utils/function_calling";
 import { AnthropicToolsOutputParser } from "./output_parsers.js";
-import { extractToolCallChunk, handleToolChoice } from "./utils/tools.js";
+import { handleToolChoice } from "./utils/tools.js";
 import { _convertMessagesToAnthropicPayload } from "./utils/message_inputs.js";
 import {
   _makeMessageChunkFromAnthropicEvent,
@@ -815,8 +815,6 @@ export class ChatAnthropicMessages<
 
       const { chunk } = result;
 
-      const newToolCallChunk = extractToolCallChunk(chunk);
-
       // Extract the text content token for text field and runManager.
       const token = extractToken(chunk);
       const generationChunk = new ChatGenerationChunk({
@@ -824,7 +822,7 @@ export class ChatAnthropicMessages<
           // Just yield chunk as it is and tool_use will be concat by BaseChatModel._generateUncached().
           content: chunk.content,
           additional_kwargs: chunk.additional_kwargs,
-          tool_call_chunks: newToolCallChunk ? [newToolCallChunk] : undefined,
+          tool_call_chunks: chunk.tool_call_chunks,
           usage_metadata: shouldStreamUsage ? chunk.usage_metadata : undefined,
           response_metadata: chunk.response_metadata,
           id: chunk.id,
diff --git a/libs/langchain-anthropic/src/utils/message_outputs.ts b/libs/langchain-anthropic/src/utils/message_outputs.ts
index a04dfba44110..fd34dba87bf7 100644
--- a/libs/langchain-anthropic/src/utils/message_outputs.ts
+++ b/libs/langchain-anthropic/src/utils/message_outputs.ts
@@ -60,6 +60,8 @@ export function _makeMessageChunkFromAnthropicEvent(
     data.type === "content_block_start" &&
     data.content_block.type === "tool_use"
   ) {
+    const toolCallContentBlock =
+      data.content_block as Anthropic.Messages.ToolUseBlock;
     return {
       chunk: new AIMessageChunk({
         content: fields.coerceContentToString
@@ -72,6 +74,14 @@ export function _makeMessageChunkFromAnthropicEvent(
               },
             ],
         additional_kwargs: {},
+        tool_call_chunks: [
+          {
+            id: toolCallContentBlock.id,
+            index: data.index,
+            name: toolCallContentBlock.name,
+            args: "",
+          },
+        ],
       }),
     };
   } else if (
@@ -110,6 +120,12 @@ export function _makeMessageChunkFromAnthropicEvent(
               },
             ],
         additional_kwargs: {},
+        tool_call_chunks: [
+          {
+            index: data.index,
+            args: data.delta.partial_json,
+          },
+        ],
       }),
     };
   } else if (
diff --git a/libs/langchain-anthropic/src/utils/tools.ts b/libs/langchain-anthropic/src/utils/tools.ts
index a56bd22e2a49..157caa1b7f11 100644
--- a/libs/langchain-anthropic/src/utils/tools.ts
+++ b/libs/langchain-anthropic/src/utils/tools.ts
@@ -1,6 +1,4 @@
 import type { MessageCreateParams } from "@anthropic-ai/sdk/resources/index.mjs";
-import { AIMessageChunk } from "@langchain/core/messages";
-import { ToolCallChunk } from "@langchain/core/messages/tool";
 import { AnthropicToolChoice } from "../types.js";
 
 export function handleToolChoice(
@@ -29,59 +27,3 @@ export function handleToolChoice(
     return toolChoice;
   }
 }
-
-export function extractToolCallChunk(
-  chunk: AIMessageChunk
-): ToolCallChunk | undefined {
-  let newToolCallChunk: ToolCallChunk | undefined;
-
-  // Initial chunk for tool calls from anthropic contains identifying information like ID and name.
-  // This chunk does not contain any input JSON.
-  const toolUseChunks = Array.isArray(chunk.content)
-    ? chunk.content.find((c) => c.type === "tool_use")
-    : undefined;
-  if (
-    toolUseChunks &&
-    "index" in toolUseChunks &&
-    "name" in toolUseChunks &&
-    "id" in toolUseChunks
-  ) {
-    newToolCallChunk = {
-      args: "",
-      id: toolUseChunks.id,
-      name: toolUseChunks.name,
-      index: toolUseChunks.index,
-      type: "tool_call_chunk",
-    };
-  }
-
-  // Chunks after the initial chunk only contain the index and partial JSON.
-  const inputJsonDeltaChunks = Array.isArray(chunk.content)
-    ? chunk.content.find((c) => c.type === "input_json_delta")
-    : undefined;
-  if (
-    inputJsonDeltaChunks &&
-    "index" in inputJsonDeltaChunks &&
-    "input" in inputJsonDeltaChunks
-  ) {
-    if (typeof inputJsonDeltaChunks.input === "string") {
-      newToolCallChunk = {
-        id: inputJsonDeltaChunks.id,
-        name: inputJsonDeltaChunks.name,
-        args: inputJsonDeltaChunks.input,
-        index: inputJsonDeltaChunks.index,
-        type: "tool_call_chunk",
-      };
-    } else {
-      newToolCallChunk = {
-        id: inputJsonDeltaChunks.id,
-        name: inputJsonDeltaChunks.name,
-        args: JSON.stringify(inputJsonDeltaChunks.input, null, 2),
-        index: inputJsonDeltaChunks.index,
-        type: "tool_call_chunk",
-      };
-    }
-  }
-
-  return newToolCallChunk;
-}