mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(transformations): raise on invalid completion responses for fireworks and ollama
This commit is contained in:
Mohammed Alshyakh
parent
16645cc963
commit
f82c6072c2
4 changed files with 179 additions and 170 deletions
|
|
@ -760,6 +760,20 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
|
|||
additional_headers: Final = get_response_headers(raw_response_headers)
|
||||
|
||||
completion_fields: Final = _JSON_DICT.validate_python(completion_response)
|
||||
if "error" in completion_fields:
|
||||
raise FireworksAIException(
|
||||
message=f"Fireworks AI returned an error body with status {raw_response.status_code}: {completion_fields['error']}",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response_headers,
|
||||
)
|
||||
|
||||
if not isinstance(completion_fields.get("choices"), list):
|
||||
raise FireworksAIException(
|
||||
message=f"LiteLLM: provider returned a response with no 'choices'. Raw keys: {list(completion_fields.keys())}",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response_headers,
|
||||
)
|
||||
|
||||
response: Final = ModelResponse(**completion_fields)
|
||||
|
||||
if response.model is not None:
|
||||
|
|
|
|||
|
|
@ -321,6 +321,20 @@ class OllamaConfig(BaseConfig):
|
|||
json_mode: bool | None = None,
|
||||
) -> ModelResponse:
|
||||
response_json: Final = _OLLAMA_GENERATE_RESPONSE.validate_python(raw_response.json())
|
||||
if "error" in response_json:
|
||||
raise OllamaError(
|
||||
status_code=raw_response.status_code,
|
||||
message=f"Ollama returned an error body with status {raw_response.status_code}: {response_json['error']}",
|
||||
headers=dict(raw_response.headers),
|
||||
)
|
||||
|
||||
if response_json.get("response") is None and response_json.get("thinking") is None:
|
||||
raw_keys: Final = list(response_json.keys())
|
||||
raise OllamaError(
|
||||
status_code=raw_response.status_code,
|
||||
message=f"LiteLLM: provider returned a response with no 'response'. Raw keys: {raw_keys}",
|
||||
headers=dict(raw_response.headers),
|
||||
)
|
||||
## RESPONSE OBJECT
|
||||
model_response.choices[0].finish_reason = "stop"
|
||||
if request_data.get("format", "") == "json":
|
||||
|
|
|
|||
|
|
@ -12,7 +12,10 @@ from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
|
|||
from litellm.litellm_core_utils.get_supported_openai_params import get_supported_openai_params
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
|
||||
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
|
||||
from litellm.llms.fireworks_ai.common_utils import (
|
||||
FireworksAIException,
|
||||
get_fireworks_session_id,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionMessageToolCall,
|
||||
Function,
|
||||
|
|
@ -221,9 +224,7 @@ def test_validate_environment_raises_without_api_key(monkeypatch):
|
|||
|
||||
def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id():
|
||||
assert (
|
||||
get_fireworks_session_id(
|
||||
{"litellm_session_id": "session-123", "litellm_trace_id": "trace-123"}
|
||||
)
|
||||
get_fireworks_session_id({"litellm_session_id": "session-123", "litellm_trace_id": "trace-123"})
|
||||
== "session-123"
|
||||
)
|
||||
|
||||
|
|
@ -275,25 +276,18 @@ def test_handle_message_content_with_tool_calls():
|
|||
},
|
||||
}
|
||||
]
|
||||
updated_message = config._handle_message_content_with_tool_calls(
|
||||
message, tool_calls
|
||||
)
|
||||
updated_message = config._handle_message_content_with_tool_calls(message, tool_calls)
|
||||
assert updated_message.tool_calls is not None
|
||||
assert len(updated_message.tool_calls) == 1
|
||||
assert updated_message.tool_calls[0].function.name == "get_current_weather"
|
||||
assert (
|
||||
updated_message.tool_calls[0].function.arguments
|
||||
== expected_tool_call.function.arguments
|
||||
)
|
||||
assert updated_message.tool_calls[0].function.arguments == expected_tool_call.function.arguments
|
||||
|
||||
|
||||
def test_get_supported_openai_params_reasoning_effort():
|
||||
"""Test that reasoning_effort is only included in supported params for models that support it."""
|
||||
config = FireworksAIConfig()
|
||||
|
||||
supported_params = config.get_supported_openai_params(
|
||||
"fireworks_ai/accounts/fireworks/models/glm-5p1"
|
||||
)
|
||||
supported_params = config.get_supported_openai_params("fireworks_ai/accounts/fireworks/models/glm-5p1")
|
||||
assert "reasoning_effort" in supported_params
|
||||
assert "thinking" in supported_params
|
||||
|
||||
|
|
@ -308,9 +302,7 @@ def test_get_supported_openai_params_parallel_tool_calls():
|
|||
"""Test that parallel_tool_calls is included for models that support function calling."""
|
||||
config = FireworksAIConfig()
|
||||
|
||||
supported_params = config.get_supported_openai_params(
|
||||
"fireworks_ai/accounts/fireworks/models/glm-5p1"
|
||||
)
|
||||
supported_params = config.get_supported_openai_params("fireworks_ai/accounts/fireworks/models/glm-5p1")
|
||||
assert "parallel_tool_calls" in supported_params
|
||||
assert "tools" in supported_params
|
||||
assert "tool_choice" in supported_params
|
||||
|
|
@ -324,9 +316,7 @@ def test_get_supported_openai_params_parallel_tool_calls():
|
|||
def test_get_supported_openai_params_short_model_name_resolves_account_prefixed_entry():
|
||||
config = FireworksAIConfig()
|
||||
|
||||
supported_params = config.get_supported_openai_params(
|
||||
"fireworks_ai/deepseek-v4-pro-0813"
|
||||
)
|
||||
supported_params = config.get_supported_openai_params("fireworks_ai/deepseek-v4-pro-0813")
|
||||
|
||||
assert "tool_choice" in supported_params
|
||||
assert "reasoning_effort" in supported_params
|
||||
|
|
@ -335,9 +325,7 @@ def test_get_supported_openai_params_short_model_name_resolves_account_prefixed_
|
|||
def test_get_supported_openai_params_preserves_generic_reasoning_fallback():
|
||||
config = FireworksAIConfig()
|
||||
|
||||
supported_params = config.get_supported_openai_params(
|
||||
"fireworks_ai/accounts/fireworks/models/glm-5p3-flash"
|
||||
)
|
||||
supported_params = config.get_supported_openai_params("fireworks_ai/accounts/fireworks/models/glm-5p3-flash")
|
||||
|
||||
assert "reasoning_effort" in supported_params
|
||||
|
||||
|
|
@ -404,14 +392,10 @@ def test_get_models_url_no_double_v1(api_base, expected_url_prefix):
|
|||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"models": [{"name": "accounts/fireworks/models/llama-v3-70b"}]
|
||||
}
|
||||
mock_response.json.return_value = {"models": [{"name": "accounts/fireworks/models/llama-v3-70b"}]}
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.module_level_client.get", return_value=mock_response
|
||||
) as mock_get,
|
||||
patch("litellm.module_level_client.get", return_value=mock_response) as mock_get,
|
||||
patch(
|
||||
"litellm.llms.fireworks_ai.chat.transformation.get_secret_str",
|
||||
side_effect=lambda key: {
|
||||
|
|
@ -423,13 +407,9 @@ def test_get_models_url_no_double_v1(api_base, expected_url_prefix):
|
|||
):
|
||||
result = config.get_models(api_key="test-key", api_base=api_base)
|
||||
|
||||
called_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args[1].get(
|
||||
"url", ""
|
||||
)
|
||||
called_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args[1].get("url", "")
|
||||
assert "/v1/v1/" not in called_url, f"Double /v1/ detected in URL: {called_url}"
|
||||
assert called_url.startswith(
|
||||
expected_url_prefix
|
||||
), f"URL {called_url} does not start with {expected_url_prefix}"
|
||||
assert called_url.startswith(expected_url_prefix), f"URL {called_url} does not start with {expected_url_prefix}"
|
||||
assert result == ["fireworks_ai/accounts/fireworks/models/llama-v3-70b"]
|
||||
|
||||
|
||||
|
|
@ -457,9 +437,7 @@ def test_transform_messages_helper_removes_provider_specific_fields():
|
|||
},
|
||||
]
|
||||
# Call helper
|
||||
out = config._transform_messages_helper(
|
||||
messages, model="fireworks/test", litellm_params={}
|
||||
)
|
||||
out = config._transform_messages_helper(messages, model="fireworks/test", litellm_params={})
|
||||
for msg in out:
|
||||
assert "provider_specific_fields" not in msg
|
||||
|
||||
|
|
@ -480,15 +458,11 @@ def test_transform_messages_helper_strips_thinking_blocks_but_keeps_reasoning_co
|
|||
{
|
||||
"role": "assistant",
|
||||
"content": "I can help.",
|
||||
"thinking_blocks": [
|
||||
{"type": "thinking", "thinking": "internal", "signature": ""}
|
||||
],
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": "internal", "signature": ""}],
|
||||
"reasoning_content": "internal",
|
||||
},
|
||||
]
|
||||
out = config._transform_messages_helper(
|
||||
messages, model="accounts/fireworks/models/glm-5p1", litellm_params={}
|
||||
)
|
||||
out = config._transform_messages_helper(messages, model="accounts/fireworks/models/glm-5p1", litellm_params={})
|
||||
assert "thinking_blocks" not in out[1]
|
||||
assert out[1]["reasoning_content"] == "internal"
|
||||
assert out[1]["content"] == "I can help."
|
||||
|
|
@ -1006,9 +980,7 @@ def test_transform_messages_helper_rejects_file_blocks():
|
|||
litellm.BadRequestError,
|
||||
match="Fireworks AI chat completions does not support file content blocks",
|
||||
):
|
||||
config._transform_messages_helper(
|
||||
messages, model="accounts/fireworks/models/kimi-k2p6", litellm_params={}
|
||||
)
|
||||
config._transform_messages_helper(messages, model="accounts/fireworks/models/kimi-k2p6", litellm_params={})
|
||||
|
||||
|
||||
def test_transform_messages_helper_rejects_non_vision_image_inputs():
|
||||
|
|
@ -1020,18 +992,14 @@ def test_transform_messages_helper_rejects_non_vision_image_inputs():
|
|||
{"type": "text", "text": "Describe this"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE="
|
||||
},
|
||||
"image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE="},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match="does not support image inputs"):
|
||||
config._transform_messages_helper(
|
||||
messages, model="accounts/fireworks/models/glm-5p2", litellm_params={}
|
||||
)
|
||||
config._transform_messages_helper(messages, model="accounts/fireworks/models/glm-5p2", litellm_params={})
|
||||
|
||||
|
||||
def test_transform_messages_helper_allows_vision_image_inputs():
|
||||
|
|
@ -1043,9 +1011,7 @@ def test_transform_messages_helper_allows_vision_image_inputs():
|
|||
{"type": "text", "text": "Describe this"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE="
|
||||
},
|
||||
"image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE="},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
|
@ -1069,9 +1035,7 @@ def test_image_inputs_not_rejected_for_fuzzy_non_vision_match():
|
|||
custom_model = "accounts/myorg/models/custom-glm-5p2"
|
||||
|
||||
assert config._get_model_cost_capability(custom_model, "supports_vision") is False
|
||||
assert (
|
||||
config._get_model_cost_capability_exact(custom_model, "supports_vision") is None
|
||||
)
|
||||
assert config._get_model_cost_capability_exact(custom_model, "supports_vision") is None
|
||||
|
||||
messages = [
|
||||
{
|
||||
|
|
@ -1079,16 +1043,12 @@ def test_image_inputs_not_rejected_for_fuzzy_non_vision_match():
|
|||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE="
|
||||
},
|
||||
"image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAE="},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
out = config._transform_messages_helper(
|
||||
messages, model=custom_model, litellm_params={}
|
||||
)
|
||||
out = config._transform_messages_helper(messages, model=custom_model, litellm_params={})
|
||||
assert out == messages
|
||||
|
||||
|
||||
|
|
@ -1101,9 +1061,7 @@ def test_transform_messages_helper_skips_non_dict_content():
|
|||
}
|
||||
]
|
||||
|
||||
out = config._transform_messages_helper(
|
||||
messages, model="accounts/fireworks/models/glm-5p2", litellm_params={}
|
||||
)
|
||||
out = config._transform_messages_helper(messages, model="accounts/fireworks/models/glm-5p2", litellm_params={})
|
||||
assert out == messages
|
||||
|
||||
|
||||
|
|
@ -1343,9 +1301,7 @@ def test_streaming_surfaces_fireworks_response_fields():
|
|||
surfaced: dict = {}
|
||||
for chunk in stream:
|
||||
fields = getattr(chunk, "provider_specific_fields", None) or {}
|
||||
surfaced.update(
|
||||
{k: v for k, v in fields.items() if k.startswith("fireworks_")}
|
||||
)
|
||||
surfaced.update({k: v for k, v in fields.items() if k.startswith("fireworks_")})
|
||||
|
||||
assert surfaced["fireworks_token_ids"] == [[123]]
|
||||
assert surfaced["fireworks_raw_outputs"] == [raw_output]
|
||||
|
|
@ -1398,9 +1354,7 @@ def test_transform_request_direct_route_passthrough():
|
|||
|
||||
def test_map_extra_body_params_translates_truncate_prompt_tokens():
|
||||
config = FireworksAIConfig()
|
||||
result = config.map_extra_body_params(
|
||||
{"extra_body": {"truncate_prompt_tokens": 4096}}, _REASONING_MODEL
|
||||
)
|
||||
result = config.map_extra_body_params({"extra_body": {"truncate_prompt_tokens": 4096}}, _REASONING_MODEL)
|
||||
assert result == {"prompt_truncate_len": 4096}
|
||||
|
||||
|
||||
|
|
@ -1559,9 +1513,7 @@ def test_map_extra_body_params_non_dict_chat_template_kwargs_dropped():
|
|||
def test_map_extra_body_params_guided_json():
|
||||
config = FireworksAIConfig()
|
||||
schema = {"type": "object", "properties": {"x": {"type": "string"}}}
|
||||
result = config.map_extra_body_params(
|
||||
{"extra_body": {"guided_json": schema}}, _REASONING_MODEL
|
||||
)
|
||||
result = config.map_extra_body_params({"extra_body": {"guided_json": schema}}, _REASONING_MODEL)
|
||||
assert result == {
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
|
|
@ -1572,16 +1524,10 @@ def test_map_extra_body_params_guided_json():
|
|||
|
||||
def test_map_extra_body_params_guided_grammar_and_choice():
|
||||
config = FireworksAIConfig()
|
||||
grammar = config.map_extra_body_params(
|
||||
{"extra_body": {"guided_grammar": "root ::= 'hello'"}}, _REASONING_MODEL
|
||||
)
|
||||
assert grammar == {
|
||||
"response_format": {"type": "grammar", "grammar": "root ::= 'hello'"}
|
||||
}
|
||||
grammar = config.map_extra_body_params({"extra_body": {"guided_grammar": "root ::= 'hello'"}}, _REASONING_MODEL)
|
||||
assert grammar == {"response_format": {"type": "grammar", "grammar": "root ::= 'hello'"}}
|
||||
|
||||
choice = config.map_extra_body_params(
|
||||
{"extra_body": {"guided_choice": ["yes", "no"]}}, _REASONING_MODEL
|
||||
)
|
||||
choice = config.map_extra_body_params({"extra_body": {"guided_choice": ["yes", "no"]}}, _REASONING_MODEL)
|
||||
assert choice == {
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
|
|
@ -1667,9 +1613,7 @@ def test_map_extra_body_params_strips_unsupported_nim_vllm_params(param, value,
|
|||
|
||||
config = FireworksAIConfig()
|
||||
with caplog.at_level(logging.DEBUG):
|
||||
result = config.map_extra_body_params(
|
||||
{"extra_body": {param: value}}, _REASONING_MODEL
|
||||
)
|
||||
result = config.map_extra_body_params({"extra_body": {param: value}}, _REASONING_MODEL)
|
||||
assert result == {}
|
||||
assert param in caplog.text
|
||||
|
||||
|
|
@ -1741,15 +1685,10 @@ def test_nim_vllm_extras_translated_end_to_end_in_request_body():
|
|||
assert request_body["top_k"] == 40
|
||||
|
||||
|
||||
|
||||
|
||||
def test_streaming_preserves_selected_model_for_private_accounting():
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
requested_route = (
|
||||
"accounts/fireworks/routers/firerouter/"
|
||||
"kimi-k3/deepseek-v4-pro-0813/deepseek-v4-flash-0731"
|
||||
)
|
||||
requested_route = "accounts/fireworks/routers/firerouter/kimi-k3/deepseek-v4-pro-0813/deepseek-v4-flash-0731"
|
||||
selected_model = "deepseek-v4-flash-0731"
|
||||
sse_lines = [
|
||||
"data: "
|
||||
|
|
@ -1803,19 +1742,14 @@ def test_streaming_preserves_selected_model_for_private_accounting():
|
|||
|
||||
assert chunks
|
||||
assert {chunk.model for chunk in chunks} == {requested_route}
|
||||
assert {
|
||||
chunk._hidden_params.get("provider_response_model") for chunk in chunks
|
||||
} == {selected_model}
|
||||
assert {chunk._hidden_params.get("provider_response_model") for chunk in chunks} == {selected_model}
|
||||
|
||||
assembled = litellm.stream_chunk_builder(chunks=chunks)
|
||||
assert assembled is not None
|
||||
assert assembled.model == requested_route
|
||||
assert assembled._hidden_params["provider_response_model"] == selected_model
|
||||
selected_model_info = litellm.model_cost[f"fireworks_ai/{selected_model}"]
|
||||
expected_cost = (
|
||||
5 * selected_model_info["input_cost_per_token"]
|
||||
+ selected_model_info["output_cost_per_token"]
|
||||
)
|
||||
expected_cost = 5 * selected_model_info["input_cost_per_token"] + selected_model_info["output_cost_per_token"]
|
||||
assert litellm.completion_cost(
|
||||
completion_response=assembled,
|
||||
custom_llm_provider="fireworks_ai",
|
||||
|
|
@ -1914,27 +1848,19 @@ def test_listed_router_request_is_sent_to_the_router_resource_and_billed_at_the_
|
|||
|
||||
def test_map_openai_params_tool_choice():
|
||||
# Test case 1: tool_choice is "required"
|
||||
result = fireworks.map_openai_params(
|
||||
{"tool_choice": "required"}, {}, "some_model", drop_params=False
|
||||
)
|
||||
result = fireworks.map_openai_params({"tool_choice": "required"}, {}, "some_model", drop_params=False)
|
||||
assert result == {"tool_choice": "any"}
|
||||
|
||||
# Test case 2: tool_choice is "auto"
|
||||
result = fireworks.map_openai_params(
|
||||
{"tool_choice": "auto"}, {}, "some_model", drop_params=False
|
||||
)
|
||||
result = fireworks.map_openai_params({"tool_choice": "auto"}, {}, "some_model", drop_params=False)
|
||||
assert result == {"tool_choice": "auto"}
|
||||
|
||||
# Test case 3: tool_choice is not present
|
||||
result = fireworks.map_openai_params(
|
||||
{"some_other_param": "value"}, {}, "some_model", drop_params=False
|
||||
)
|
||||
result = fireworks.map_openai_params({"some_other_param": "value"}, {}, "some_model", drop_params=False)
|
||||
assert result == {}
|
||||
|
||||
# Test case 4: tool_choice is None
|
||||
result = fireworks.map_openai_params(
|
||||
{"tool_choice": None}, {}, "some_model", drop_params=False
|
||||
)
|
||||
result = fireworks.map_openai_params({"tool_choice": None}, {}, "some_model", drop_params=False)
|
||||
assert result == {"tool_choice": None}
|
||||
|
||||
|
||||
|
|
@ -1961,9 +1887,7 @@ def test_map_response_format():
|
|||
"strict": True,
|
||||
},
|
||||
}
|
||||
result = fireworks.map_openai_params(
|
||||
{"response_format": response_format}, {}, "some_model", drop_params=False
|
||||
)
|
||||
result = fireworks.map_openai_params({"response_format": response_format}, {}, "some_model", drop_params=False)
|
||||
assert result == {"response_format": response_format}
|
||||
|
||||
|
||||
|
|
@ -2358,3 +2282,35 @@ def test_global_disable_flag_with_transform_messages_helper(
|
|||
)
|
||||
|
||||
assert json.loads(route.calls.last.request.content)["messages"] == [{"role": "user", "content": content}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"invalid_body",
|
||||
[
|
||||
{},
|
||||
{"choices": None},
|
||||
{"choices": "not-a-list"},
|
||||
{"error": {"message": "upstream broke"}},
|
||||
{"error": "bad request"},
|
||||
],
|
||||
)
|
||||
def test_fireworks_invalid_body_raises(invalid_body: dict) -> None:
|
||||
config = FireworksAIConfig()
|
||||
raw_response = MagicMock()
|
||||
raw_response.json.return_value = invalid_body
|
||||
raw_response.status_code = 200
|
||||
raw_response.headers = {}
|
||||
model_response = ModelResponse()
|
||||
|
||||
with pytest.raises(FireworksAIException):
|
||||
config.transform_response(
|
||||
model="accounts/fireworks/models/glm-5p1",
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import respx
|
|||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.ollama.common_utils import OllamaError
|
||||
from litellm.llms.ollama.completion.transformation import (
|
||||
OllamaConfig,
|
||||
OllamaTextCompletionResponseIterator,
|
||||
|
|
@ -76,9 +77,7 @@ class TestOllamaConfig:
|
|||
# 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"}}
|
||||
)
|
||||
"response": json.dumps({"name": "get_weather", "arguments": {"location": "San Francisco"}})
|
||||
}
|
||||
|
||||
# Create properly structured model response object
|
||||
|
|
@ -109,13 +108,10 @@ class TestOllamaConfig:
|
|||
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"}
|
||||
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_json_function_call_in_code_fence(self):
|
||||
|
|
@ -147,9 +143,7 @@ class TestOllamaConfig:
|
|||
# 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"}}
|
||||
)
|
||||
"response": json.dumps({"result": "success", "data": {"temperature": 72, "unit": "F"}})
|
||||
}
|
||||
|
||||
# Create properly structured model response object
|
||||
|
|
@ -176,9 +170,7 @@ class TestOllamaConfig:
|
|||
)
|
||||
|
||||
# Verify result has JSON content
|
||||
expected_content = json.dumps(
|
||||
{"result": "success", "data": {"temperature": 72, "unit": "F"}}
|
||||
)
|
||||
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
|
||||
|
|
@ -220,10 +212,7 @@ class TestOllamaConfig:
|
|||
)
|
||||
|
||||
# Verify reasoning content is extracted
|
||||
assert (
|
||||
result.choices[0]["message"].reasoning_content
|
||||
== "I need to think about this problem step by step"
|
||||
)
|
||||
assert result.choices[0]["message"].reasoning_content == "I need to think about this problem step by step"
|
||||
assert result.choices[0]["message"].content == "Here is my answer"
|
||||
assert result.choices[0]["finish_reason"] == "stop"
|
||||
|
||||
|
|
@ -262,10 +251,7 @@ class TestOllamaConfig:
|
|||
)
|
||||
|
||||
# Verify reasoning content is extracted
|
||||
assert (
|
||||
result.choices[0]["message"].reasoning_content
|
||||
== "Let me analyze this carefully"
|
||||
)
|
||||
assert result.choices[0]["message"].reasoning_content == "Let me analyze this carefully"
|
||||
assert result.choices[0]["message"].content == "The solution is X"
|
||||
assert result.choices[0]["finish_reason"] == "stop"
|
||||
|
||||
|
|
@ -306,10 +292,7 @@ class TestOllamaConfig:
|
|||
# Verify multiline reasoning content is extracted
|
||||
expected_reasoning = "\nThis is a complex problem.\nI need to break it down:\n1. First step\n2. Second step\n"
|
||||
assert result.choices[0]["message"].reasoning_content == expected_reasoning
|
||||
assert (
|
||||
result.choices[0]["message"].content
|
||||
== "Based on my analysis, the answer is Y"
|
||||
)
|
||||
assert result.choices[0]["message"].content == "Based on my analysis, the answer is Y"
|
||||
assert result.choices[0]["finish_reason"] == "stop"
|
||||
|
||||
def test_transform_response_thinking_only(self):
|
||||
|
|
@ -347,10 +330,7 @@ class TestOllamaConfig:
|
|||
)
|
||||
|
||||
# Verify reasoning content is extracted and content is empty
|
||||
assert (
|
||||
result.choices[0]["message"].reasoning_content
|
||||
== "Just internal thoughts, no response"
|
||||
)
|
||||
assert result.choices[0]["message"].reasoning_content == "Just internal thoughts, no response"
|
||||
assert result.choices[0]["message"].content == ""
|
||||
assert result.choices[0]["finish_reason"] == "stop"
|
||||
|
||||
|
|
@ -389,10 +369,7 @@ class TestOllamaConfig:
|
|||
)
|
||||
|
||||
# Verify reasoning content is extracted even in JSON mode when JSON parsing fails
|
||||
assert (
|
||||
result.choices[0]["message"].reasoning_content
|
||||
== "Planning my JSON response"
|
||||
)
|
||||
assert result.choices[0]["message"].reasoning_content == "Planning my JSON response"
|
||||
assert result.choices[0]["message"].content == "This is not valid JSON"
|
||||
assert result.choices[0]["finish_reason"] == "stop"
|
||||
|
||||
|
|
@ -432,10 +409,7 @@ class TestOllamaConfig:
|
|||
|
||||
# Verify no reasoning content is extracted
|
||||
assert result.choices[0]["message"].reasoning_content is None
|
||||
assert (
|
||||
result.choices[0]["message"].content
|
||||
== "Regular response without any thinking tags"
|
||||
)
|
||||
assert result.choices[0]["message"].content == "Regular response without any thinking tags"
|
||||
assert result.choices[0]["finish_reason"] == "stop"
|
||||
|
||||
def _transform(
|
||||
|
|
@ -603,7 +577,13 @@ class TestOllamaTextCompletionResponseIterator:
|
|||
)
|
||||
ollama_chunks: Final = (
|
||||
{"model": "qwen3:0.6b", "created_at": "2026-10-07T00:00:00Z", "response": "", "done": False},
|
||||
{"model": "qwen3:0.6b", "created_at": "2026-10-07T00:00:00Z", "response": "", "thinking": "Hm", "done": False},
|
||||
{
|
||||
"model": "qwen3:0.6b",
|
||||
"created_at": "2026-10-07T00:00:00Z",
|
||||
"response": "",
|
||||
"thinking": "Hm",
|
||||
"done": False,
|
||||
},
|
||||
{"model": "qwen3:0.6b", "created_at": "2026-10-07T00:00:00Z", "response": "Hel", "done": False},
|
||||
{"model": "qwen3:0.6b", "created_at": "2026-10-07T00:00:00Z", "response": "lo", "done": False},
|
||||
)
|
||||
|
|
@ -621,9 +601,7 @@ class TestOllamaTextCompletionResponseIterator:
|
|||
|
||||
def test_chunk_parser_with_thinking_field(self):
|
||||
"""Test that chunks with 'thinking' field and empty 'response' are handled correctly."""
|
||||
iterator = OllamaTextCompletionResponseIterator(
|
||||
streaming_response=iter([]), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False)
|
||||
|
||||
# Test chunk with thinking field - this is the problematic case from the issue
|
||||
chunk_with_thinking = {
|
||||
|
|
@ -643,9 +621,7 @@ class TestOllamaTextCompletionResponseIterator:
|
|||
|
||||
def test_chunk_parser_normal_response(self):
|
||||
"""Test that normal response chunks still work."""
|
||||
iterator = OllamaTextCompletionResponseIterator(
|
||||
streaming_response=iter([]), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False)
|
||||
|
||||
# Test normal chunk with response
|
||||
normal_chunk = {
|
||||
|
|
@ -665,9 +641,7 @@ class TestOllamaTextCompletionResponseIterator:
|
|||
|
||||
def test_chunk_parser_empty_response_without_thinking(self):
|
||||
"""Test that empty response chunks without thinking still work."""
|
||||
iterator = OllamaTextCompletionResponseIterator(
|
||||
streaming_response=iter([]), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False)
|
||||
|
||||
# Test empty response chunk without thinking
|
||||
empty_response_chunk = {
|
||||
|
|
@ -687,9 +661,7 @@ class TestOllamaTextCompletionResponseIterator:
|
|||
|
||||
def test_chunk_parser_done_chunk(self):
|
||||
"""Test that done chunks work correctly."""
|
||||
iterator = OllamaTextCompletionResponseIterator(
|
||||
streaming_response=iter([]), sync_stream=True, json_mode=False
|
||||
)
|
||||
iterator = OllamaTextCompletionResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False)
|
||||
|
||||
# Test done chunk
|
||||
done_chunk = {
|
||||
|
|
@ -990,3 +962,56 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str):
|
|||
)
|
||||
== "http://gateway.example.com/ollama/api/generate"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"invalid_body",
|
||||
[
|
||||
{},
|
||||
{"choices": None},
|
||||
{"error": {"message": "upstream broke"}},
|
||||
{"error": "model not found"},
|
||||
],
|
||||
)
|
||||
def test_ollama_invalid_body_raises(invalid_body: dict) -> None:
|
||||
config = OllamaConfig()
|
||||
raw_response = MagicMock()
|
||||
raw_response.json.return_value = invalid_body
|
||||
raw_response.status_code = 200
|
||||
raw_response.headers = {}
|
||||
model_response = ModelResponse(id="test", choices=[{"message": Message(content="")}])
|
||||
|
||||
with pytest.raises(OllamaError):
|
||||
config.transform_response(
|
||||
model="llama3",
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
|
||||
def test_ollama_valid_empty_response_allowed() -> None:
|
||||
config = OllamaConfig()
|
||||
raw_response = MagicMock()
|
||||
raw_response.json.return_value = {"model": "llama3", "response": "", "done": True}
|
||||
raw_response.status_code = 200
|
||||
raw_response.headers = {}
|
||||
model_response = ModelResponse(id="test", choices=[{"message": Message(content="")}])
|
||||
|
||||
res = config.transform_response(
|
||||
model="llama3",
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
assert res.choices[0].message.content == ""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue