update cohere to use tool call attribute on AIMessage

libs/cohere/langchain_cohere/chat_models.py

@@ -32,6 +32,9 @@
     HumanMessage,
     SystemMessage,
 )
+from langchain_core.messages import (
+    ToolCall as LC_ToolCall,
+)
 from langchain_core.output_parsers.base import OutputParserLike
 from langchain_core.output_parsers.openai_tools import (
     JsonOutputKeyToolsParser,
@@ -250,10 +253,27 @@ def _stream(
                 yield chunk
             elif data.event_type == "stream-end":
                 generation_info = self._get_generation_info(data.response)
+                tool_call_chunks = []
+                if tool_calls := generation_info.get("tool_calls"):
+                    try:
+                        tool_call_chunks = [
+                            {
+                                "name": tool_call["function"].get("name"),
+                                "args": tool_call["function"].get("arguments"),
+                                "id": tool_call.get("id"),
+                                "index": tool_call.get("index"),
+                            }
+                            for tool_call in tool_calls
+                        ]
+                    except KeyError:
+                        pass
+                message = AIMessageChunk(
+                    content="",
+                    additional_kwargs=generation_info,
+                    tool_call_chunks=tool_call_chunks,
+                )
                 yield ChatGenerationChunk(
-                    message=AIMessageChunk(
-                        content="", additional_kwargs=generation_info
-                    ),
+                    message=message,
                     generation_info=generation_info,
                 )
 
@@ -282,10 +302,27 @@ async def _astream(
                 yield chunk
             elif data.event_type == "stream-end":
                 generation_info = self._get_generation_info(data.response)
+                tool_call_chunks = []
+                if tool_calls := generation_info.get("tool_calls"):
+                    try:
+                        tool_call_chunks = [
+                            {
+                                "name": tool_call["function"].get("name"),
+                                "args": tool_call["function"].get("arguments"),
+                                "id": tool_call.get("id"),
+                                "index": tool_call.get("index"),
+                            }
+                            for tool_call in tool_calls
+                        ]
+                    except KeyError:
+                        pass
+                message = AIMessageChunk(
+                    content="",
+                    additional_kwargs=generation_info,
+                    tool_call_chunks=tool_call_chunks,
+                )
                 yield ChatGenerationChunk(
-                    message=AIMessageChunk(
-                        content="", additional_kwargs=generation_info
-                    ),
+                    message=message,
                     generation_info=generation_info,
                 )
 
@@ -328,7 +365,18 @@ def _generate(
         response = self.client.chat(**request)
 
         generation_info = self._get_generation_info(response)
-        message = AIMessage(content=response.text, additional_kwargs=generation_info)
+        if "tool_calls" in generation_info:
+            tool_calls = [
+                _convert_cohere_tool_call_to_langchain(tool_call)
+                for tool_call in response.tool_calls
+            ]
+        else:
+            tool_calls = []
+        message = AIMessage(
+            content=response.text,
+            additional_kwargs=generation_info,
+            tool_calls=tool_calls,
+        )
         return ChatResult(
             generations=[
                 ChatGeneration(message=message, generation_info=generation_info)
@@ -354,7 +402,18 @@ async def _agenerate(
         response = self.client.chat(**request)
 
         generation_info = self._get_generation_info(response)
-        message = AIMessage(content=response.text, additional_kwargs=generation_info)
+        if "tool_calls" in generation_info:
+            tool_calls = [
+                _convert_cohere_tool_call_to_langchain(tool_call)
+                for tool_call in response.tool_calls
+            ]
+        else:
+            tool_calls = []
+        message = AIMessage(
+            content=response.text,
+            additional_kwargs=generation_info,
+            tool_calls=tool_calls,
+        )
         return ChatResult(
             generations=[
                 ChatGeneration(message=message, generation_info=generation_info)
@@ -388,3 +447,8 @@ def _format_cohere_tool_calls(
             }
         )
     return formatted_tool_calls
+
+
+def _convert_cohere_tool_call_to_langchain(tool_call: ToolCall) -> LC_ToolCall:
+    """Convert a Cohere tool call into langchain_core.messages.ToolCall"""
+    return LC_ToolCall(name=tool_call.name, args=tool_call.parameters, id=None)

libs/cohere/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langchain-cohere"
-version = "0.1.1"
+version = "0.1.2-rc1"
 description = "An integration package connecting Cohere and LangChain"
 authors = []
 readme = "README.md"
@@ -12,7 +12,7 @@ license = "MIT"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
-langchain-core = "^0.1.32"
+langchain-core = {version ="^0.1.42-rc.1", allow-prereleases=true}
 cohere = ">=5.1.8,<5.2"
 
 [tool.poetry.group.test]

libs/cohere/tests/integration_tests/test_chat_models.py

@@ -4,7 +4,11 @@
 from typing import Any
 
 import pytest
-from langchain_core.messages import AIMessage, AIMessageChunk
+from langchain_core.messages import (
+    AIMessage,
+    AIMessageChunk,
+    ToolCall,
+)
 from langchain_core.pydantic_v1 import BaseModel, Field
 
 from langchain_cohere import ChatCohere
@@ -92,6 +96,9 @@ class Person(BaseModel):
         "name": "Erick",
         "age": 27,
     }
+    assert result.tool_calls == [
+        ToolCall(name="Person", args={"age": 27, "name": "Erick"}, id=None)
+    ]
 
 
 def test_streaming_tool_call() -> None:
@@ -120,6 +127,13 @@ class Person(BaseModel):
         "name": "Erick",
         "age": 27,
     }
+    assert isinstance(chunk, AIMessageChunk)
+    assert isinstance(chunk.tool_call_chunks, list)
+    assert len(chunk.tool_call_chunks) == 1
+    tool_call_chunk = chunk.tool_call_chunks[0]
+    assert tool_call_chunk["name"] == "Person"
+    assert tool_call_chunk["args"] is not None
+    assert json.loads(tool_call_chunk["args"]) == {"name": "Erick", "age": 27}
 
 
 @pytest.mark.xfail(