mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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
This commit is contained in:
Aravind
GitHub
parent
ae10f3f82c
commit
4f162b0706
2 changed files with 186 additions and 16 deletions
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue