prefer .tool_calls

langchain-ai · Apr 17, 2024 · f341f4c · f341f4c
1 parent 88d4733
commit f341f4c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/libs/partners/mistralai/langchain_mistralai/chat_models.py b/libs/partners/mistralai/langchain_mistralai/chat_models.py
@@ -263,7 +263,14 @@ def _convert_message_to_mistral_chat_message(
         return dict(role="user", content=message.content)
     elif isinstance(message, AIMessage):
         tool_calls = []
-        if "tool_calls" in message.additional_kwargs:
+        if message.tool_calls or message.invalid_tool_calls:
+            for tool_call in message.tool_calls:
+                tool_calls.append(_format_tool_call_for_mistral(tool_call))
+            for invalid_tool_call in message.invalid_tool_calls:
+                tool_calls.append(
+                    _format_invalid_tool_call_for_mistral(invalid_tool_call)
+                )
+        elif "tool_calls" in message.additional_kwargs:
             for tc in message.additional_kwargs["tool_calls"]:
                 chunk = {
                     "function": {
@@ -274,13 +281,6 @@ def _convert_message_to_mistral_chat_message(
                 if _id := tc.get("id"):
                     chunk["id"] = _id
                 tool_calls.append(chunk)
-        elif message.tool_calls or message.invalid_tool_calls:
-            for tool_call in message.tool_calls:
-                tool_calls.append(_format_tool_call_for_mistral(tool_call))
-            for invalid_tool_call in message.invalid_tool_calls:
-                tool_calls.append(
-                    _format_invalid_tool_call_for_mistral(invalid_tool_call)
-                )
         else:
             pass
         return {

diff --git a/libs/partners/mistralai/tests/unit_tests/test_chat_models.py b/libs/partners/mistralai/tests/unit_tests/test_chat_models.py
@@ -130,7 +130,7 @@ def test__convert_dict_to_message_tool_call() -> None:
     raw_tool_call = {
         "id": "abc123",
         "function": {
-            "arguments": '{"name":"Sally","hair_color":"green"}',
+            "arguments": '{"name": "Sally", "hair_color": "green"}',
             "name": "GenerateUsername",
         },
     }
@@ -153,16 +153,16 @@ def test__convert_dict_to_message_tool_call() -> None:
     # Test malformed tool call
     raw_tool_calls = [
         {
-            "id": "abc123",
+            "id": "def456",
             "function": {
-                "arguments": "oops",
+                "arguments": '{"name": "Sally", "hair_color": "green"}',
                 "name": "GenerateUsername",
             },
         },
         {
-            "id": "def456",
+            "id": "abc123",
             "function": {
-                "arguments": '{"name":"Sally","hair_color":"green"}',
+                "arguments": "oops",
                 "name": "GenerateUsername",
             },
         },