diff --git a/libs/partners/ollama/langchain_ollama/chat_models.py b/libs/partners/ollama/langchain_ollama/chat_models.py index 94595a1bde5..984cc0ccfcc 100644 --- a/libs/partners/ollama/langchain_ollama/chat_models.py +++ b/libs/partners/ollama/langchain_ollama/chat_models.py @@ -125,13 +125,17 @@ def _parse_arguments_from_tool_call( if "function" not in raw_tool_call: return None arguments = raw_tool_call["function"]["arguments"] - parsed_arguments = {} + parsed_arguments: dict = {} if isinstance(arguments, dict): for key, value in arguments.items(): if isinstance(value, str): - parsed_arguments[key] = _parse_json_string( + parsed_value = _parse_json_string( value, skip=True, raw_tool_call=raw_tool_call ) + if isinstance(parsed_value, (dict, list)): + parsed_arguments[key] = parsed_value + else: + parsed_arguments[key] = value else: parsed_arguments[key] = value else: diff --git a/libs/partners/ollama/tests/unit_tests/test_chat_models.py b/libs/partners/ollama/tests/unit_tests/test_chat_models.py index 850e3423e0b..01abf25c460 100644 --- a/libs/partners/ollama/tests/unit_tests/test_chat_models.py +++ b/libs/partners/ollama/tests/unit_tests/test_chat_models.py @@ -1,10 +1,10 @@ """Test chat model integration.""" - +import json from typing import Dict, Type from langchain_tests.unit_tests import ChatModelUnitTests -from langchain_ollama.chat_models import ChatOllama +from langchain_ollama.chat_models import ChatOllama, _parse_arguments_from_tool_call class TestChatOllama(ChatModelUnitTests): @@ -15,3 +15,11 @@ class TestChatOllama(ChatModelUnitTests): @property def chat_model_params(self) -> Dict: return {"model": "llama3-groq-tool-use"} + + +def test__parse_arguments_from_tool_call() -> None: + raw_response = '{"model":"sample-model","message":{"role":"assistant","content":"","tool_calls":[{"function":{"name":"get_profile_details","arguments":{"arg_1":"12345678901234567890123456"}}}]},"done":false}' # noqa: E501 + raw_tool_calls = json.loads(raw_response)["message"]["tool_calls"] + response = _parse_arguments_from_tool_call(raw_tool_calls[0]) + assert response is not None + assert isinstance(response["arg_1"], str)