aws[patch]: Add ToolCall "type"

libs/aws/langchain_aws/chat_models/bedrock.py

@@ -18,11 +18,8 @@
 )
 
 from langchain_core._api.deprecation import deprecated
-from langchain_core.callbacks import (
-    CallbackManagerForLLMRun,
-)
-from langchain_core.language_models import LanguageModelInput
-from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.callbacks import CallbackManagerForLLMRun
+from langchain_core.language_models import BaseChatModel, LanguageModelInput
 from langchain_core.messages import (
     AIMessage,
     AIMessageChunk,
@@ -32,7 +29,7 @@
     SystemMessage,
 )
 from langchain_core.messages.ai import UsageMetadata
-from langchain_core.messages.tool import ToolCall, ToolMessage
+from langchain_core.messages.tool import ToolCall, ToolMessage, tool_call_chunk
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.pydantic_v1 import BaseModel, Extra
 from langchain_core.runnables import Runnable, RunnableMap, RunnablePassthrough
@@ -445,12 +442,12 @@ def _stream(
                 message = result.generations[0].message
                 if isinstance(message, AIMessage) and message.tool_calls is not None:
                     tool_call_chunks = [
-                        {
-                            "name": tool_call["name"],
-                            "args": json.dumps(tool_call["args"]),
-                            "id": tool_call["id"],
-                            "index": idx,
-                        }
+                        tool_call_chunk(
+                            name=tool_call["name"],
+                            args=json.dumps(tool_call["args"]),
+                            id=tool_call["id"],
+                            index=idx,
+                        )
                         for idx, tool_call in enumerate(message.tool_calls)
                     ]
                     message_chunk = AIMessageChunk(
@@ -512,7 +509,7 @@ def _generate(
             )
         completion = ""
         llm_output: Dict[str, Any] = {}
-        tool_calls: List[Dict[str, Any]] = []
+        tool_calls: List[ToolCall] = []
         provider_stop_reason_code = self.provider_stop_reason_key_map.get(
             self._get_provider(), "stop_reason"
         )

libs/aws/langchain_aws/chat_models/bedrock_converse.py

@@ -31,10 +31,11 @@
     HumanMessageChunk,
     SystemMessage,
     ToolCall,
-    ToolCallChunk,
     ToolMessage,
 )
 from langchain_core.messages.ai import AIMessageChunk, UsageMetadata
+from langchain_core.messages.tool import tool_call as create_tool_call
+from langchain_core.messages.tool import tool_call_chunk
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.pydantic_v1 import BaseModel, Extra, Field, root_validator
 from langchain_core.runnables import Runnable, RunnableMap, RunnablePassthrough
@@ -591,7 +592,7 @@ def _parse_stream_event(event: Dict[str, Any]) -> Optional[BaseMessageChunk]:
         tool_call_chunks = []
         if block["type"] == "tool_use":
             tool_call_chunks.append(
-                ToolCallChunk(
+                tool_call_chunk(
                     name=block.get("name"),
                     id=block.get("id"),
                     args=block.get("input"),
@@ -607,7 +608,7 @@ def _parse_stream_event(event: Dict[str, Any]) -> Optional[BaseMessageChunk]:
         tool_call_chunks = []
         if block["type"] == "tool_use":
             tool_call_chunks.append(
-                ToolCallChunk(
+                tool_call_chunk(
                     name=block.get("name"),
                     id=block.get("id"),
                     args=block.get("input"),
@@ -782,7 +783,9 @@ def _extract_tool_calls(anthropic_content: List[dict]) -> List[ToolCall]:
     for block in anthropic_content:
         if block["type"] == "tool_use":
             tool_calls.append(
-                ToolCall(name=block["name"], args=block["input"], id=block["id"])
+                create_tool_call(
+                    name=block["name"], args=block["input"], id=block["id"]
+                )
             )
     return tool_calls
 

libs/aws/langchain_aws/llms/bedrock.py

@@ -23,6 +23,7 @@
 )
 from langchain_core.language_models import LLM, BaseLanguageModel
 from langchain_core.messages import ToolCall
+from langchain_core.messages.tool import tool_call
 from langchain_core.outputs import Generation, GenerationChunk, LLMResult
 from langchain_core.pydantic_v1 import Extra, Field, root_validator
 from langchain_core.utils import get_from_dict_or_env
@@ -199,7 +200,7 @@ def extract_tool_calls(content: List[dict]) -> List[ToolCall]:
         if block["type"] != "tool_use":
             continue
         tool_calls.append(
-            ToolCall(name=block["name"], args=block["input"], id=block["id"])
+            tool_call(name=block["name"], args=block["input"], id=block["id"])
         )
     return tool_calls
 
@@ -632,7 +633,7 @@ def _prepare_input_and_invoke(
         **kwargs: Any,
     ) -> Tuple[
         str,
-        List[dict],
+        List[ToolCall],
         Dict[str, Any],
     ]:
         _model_kwargs = self.model_kwargs or {}