From 4f162b0706826c034c060a3081360108b4126e3d Mon Sep 17 00:00:00 2001 From: Aravind Date: Wed, 7 May 2025 10:52:04 +0530 Subject: [PATCH] Fix: Ollama integration KeyError when using JSON response format (#10611) * fix: Ollama JSON response handling to support arbitrary JSON structures. See #10589 * fix: linting issues and tests for #10589 --- .../llms/ollama/completion/transformation.py | 44 +++-- .../test_ollama_completion_transformation.py | 158 ++++++++++++++++++ 2 files changed, 186 insertions(+), 16 deletions(-) create mode 100644 tests/litellm/llms/ollama/test_ollama_completion_transformation.py diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 789b728337f..c619fd8cfb7 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -256,22 +256,34 @@ class OllamaConfig(BaseConfig): ## RESPONSE OBJECT model_response.choices[0].finish_reason = "stop" if request_data.get("format", "") == "json": - function_call = json.loads(response_json["response"]) - message = litellm.Message( - content=None, - tool_calls=[ - { - "id": f"call_{str(uuid.uuid4())}", - "function": { - "name": function_call["name"], - "arguments": json.dumps(function_call["arguments"]), - }, - "type": "function", - } - ], - ) - model_response.choices[0].message = message # type: ignore - model_response.choices[0].finish_reason = "tool_calls" + response_content = json.loads(response_json["response"]) + + # Check if this is a function call format with name/arguments structure + if isinstance(response_content, dict) and "name" in response_content and "arguments" in response_content: + # Handle as function call (original behavior) + function_call = response_content + message = litellm.Message( + content=None, + tool_calls=[ + { + "id": f"call_{str(uuid.uuid4())}", + "function": { + "name": function_call["name"], + "arguments": json.dumps(function_call["arguments"]), + }, + "type": "function", + } + ], + ) + model_response.choices[0].message = message # type: ignore + model_response.choices[0].finish_reason = "tool_calls" + else: + # Handle as regular JSON (new behavior) + message = litellm.Message( + content=json.dumps(response_content), + ) + model_response.choices[0].message = message # type: ignore + model_response.choices[0].finish_reason = "stop" else: model_response.choices[0].message.content = response_json["response"] # type: ignore model_response.created = int(time.time()) diff --git a/tests/litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/litellm/llms/ollama/test_ollama_completion_transformation.py new file mode 100644 index 00000000000..e724bc35f83 --- /dev/null +++ b/tests/litellm/llms/ollama/test_ollama_completion_transformation.py @@ -0,0 +1,158 @@ +import os +import sys +import json +import uuid +import pytest +from unittest.mock import MagicMock, patch + + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.ollama.completion.transformation import ( + OllamaConfig, +) +from litellm.types.utils import ModelResponse +from litellm.types.utils import Message + + +class TestOllamaConfig: + def test_transform_response_standard(self): + # Initialize config + config = OllamaConfig() + + # Create mock response + raw_response = MagicMock() + raw_response.json.return_value = { + "response": "Hello, I am an AI assistant", + "prompt_eval_count": 10, + "eval_count": 5 + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] # Return dummy token IDs + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify response + assert result.choices[0]["message"].content == "Hello, I am an AI assistant" + assert result.choices[0]["finish_reason"] == "stop" + assert result.model == "ollama/llama2" + assert result.created is not None + # Access usage properly + assert result["usage"]["prompt_tokens"] == 10 + assert result["usage"]["completion_tokens"] == 5 + assert result["usage"]["total_tokens"] == 15 + + @patch("uuid.uuid4") + def test_transform_response_json_function_call(self, mock_uuid4): + # Setup mock UUID + mock_uuid4.return_value = "test-uuid" + + # Initialize config + config = OllamaConfig() + + # Create mock response with JSON function call format + raw_response = MagicMock() + raw_response.json.return_value = { + "response": json.dumps({ + "name": "get_weather", + "arguments": {"location": "San Francisco"} + }) + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] # Return dummy token IDs + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={"format": "json"}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify result has tool_calls + assert result.choices[0]["message"].content is None + assert result.choices[0]["finish_reason"] == "tool_calls" + assert len(result.choices[0]["message"].tool_calls) == 1 + assert result.choices[0]["message"].tool_calls[0]["id"].startswith("call_") + assert result.choices[0]["message"].tool_calls[0]["function"]["name"] == "get_weather" + assert json.loads(result.choices[0]["message"].tool_calls[0]["function"]["arguments"]) == {"location": "San Francisco"} + # No usage assertions here as we don't need to test them in every case + + def test_transform_response_regular_json(self): + # Initialize config + config = OllamaConfig() + + # Create mock response with regular JSON (not function call) + raw_response = MagicMock() + raw_response.json.return_value = { + "response": json.dumps({ + "result": "success", + "data": {"temperature": 72, "unit": "F"} + }) + } + + # Create properly structured model response object + model_response = ModelResponse( + id="test_id", + choices=[{"message": Message(content="")}], + ) + + # Create mock encoding + mock_encoding = MagicMock() + mock_encoding.encode.return_value = [1, 2, 3] # Return dummy token IDs + + # Transform response + result = config.transform_response( + model="llama2", + raw_response=raw_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={"format": "json"}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=mock_encoding, + ) + + # Verify result has JSON content + expected_content = json.dumps({ + "result": "success", + "data": {"temperature": 72, "unit": "F"} + }) + assert result.choices[0]["message"].content == expected_content + assert result.choices[0]["finish_reason"] == "stop" + # No usage assertions here as we don't need to test them in every case \ No newline at end of file