diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index fa5794d71b3..afedf464c42 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -92,6 +92,28 @@ def _extract_fireworks_hidden_params(payload: dict) -> dict: return {**top_level, **per_choice} +def _has_fireworks_completion_payload(choice: object) -> bool: + if not isinstance(choice, dict): + return False + message: Final[object] = choice.get("message") + if not isinstance(message, dict): + return False + return any( + field in message + for field in ( + "content", + "tool_calls", + "function_call", + "refusal", + "reasoning_content", + "reasoning", + "thinking", + "audio", + "annotations", + ) + ) + + def _json_schema_response_format(schema: object, name: str) -> Mapping[str, object]: return {"type": "json_schema", "json_schema": {"name": name, "schema": schema}} @@ -761,6 +783,18 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): additional_headers: Final = get_response_headers(raw_response_headers) completion_fields: Final = _JSON_DICT.validate_python(completion_response) + choices: Final = completion_fields.get("choices") + valid_choices: Final = ( + isinstance(choices, list) + and bool(choices) + and all(_has_fireworks_completion_payload(choice) for choice in choices) + ) + if not valid_choices: + raise FireworksAIException( + message=f"Invalid Fireworks AI completion response: {raw_response.text}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) response: Final = ModelResponse(**completion_fields) if response.model is not None: diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 239b4e32645..026a706f5be 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -322,6 +322,12 @@ class OllamaConfig(BaseConfig): json_mode: bool | None = None, ) -> ModelResponse: response_json: Final = _OLLAMA_GENERATE_RESPONSE.validate_python(raw_response.json()) + if "response" not in response_json or "error" in response_json: + raise OllamaError( + status_code=raw_response.status_code, + message=f"Invalid Ollama completion response: {raw_response.text}", + headers=raw_response.headers, + ) ## RESPONSE OBJECT model_response.choices[0].finish_reason = "stop" if request_data.get("format", "") == "json": diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 04b791cf289..20e21551a32 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -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, @@ -828,6 +831,39 @@ def _run_transform_response(response_body: dict) -> ModelResponse: ) +@pytest.mark.parametrize( + "body", + [ + {}, + {"error": {"message": "upstream broke"}}, + {"choices": None}, + {"choices": "not-a-list"}, + {"choices": []}, + {"choices": [{}]}, + {"choices": [{"message": {}}]}, + ], +) +def test_transform_response_rejects_payload_without_completion_choices(body: dict): + with pytest.raises(FireworksAIException): + _run_transform_response(body) + + +def test_transform_response_accepts_empty_completion_content(): + body = { + **_BASE_CHAT_COMPLETION_RESPONSE, + "choices": [ + { + **_BASE_CHAT_COMPLETION_RESPONSE["choices"][0], + "message": {"role": "assistant", "content": ""}, + } + ], + } + + result = _run_transform_response(body) + + assert result.choices[0].message.content == "" + + _REASONING_MODEL = "fireworks_ai/accounts/fireworks/models/glm-5p1" _NON_REASONING_MODEL = "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct" @@ -2077,6 +2113,15 @@ def _complete_against_reply(reply: object) -> ModelResponse: ) +@pytest.mark.parametrize( + "reply", + [pytest.param({}, id="missing-choices"), pytest.param({"choices": [{}]}, id="malformed-choice")], +) +def test_completion_with_invalid_success_reply_raises_api_error(reply: object) -> None: + with pytest.raises(litellm.APIError): + _complete_against_reply(reply) + + def test_completion_keeps_unmodelled_reply_fields_and_nested_fireworks_metrics() -> None: perf_metrics: Final = {"ttft": 0.1, "per_stage": [{"name": "prefill", "ms": None}]} diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index 248e745ebff..26ce84fc14e 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -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, @@ -20,6 +21,34 @@ from litellm.types.utils import Message, ModelResponse, ModelResponseStream class TestOllamaConfig: + @pytest.mark.parametrize( + "body", + [ + {}, + {"error": {"message": "upstream broke"}}, + {"response": "", "error": {"message": "upstream broke"}}, + ], + ) + def test_transform_response_rejects_payload_without_completion(self, body): + raw_response = MagicMock() + raw_response.json.return_value = body + raw_response.text = json.dumps(body) + raw_response.status_code = 200 + raw_response.headers = {} + + with pytest.raises(OllamaError): + OllamaConfig().transform_response( + model="llama3", + raw_response=raw_response, + model_response=ModelResponse(choices=[{"message": Message(content="")}]), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + def test_transform_response_standard(self): # Initialize config config = OllamaConfig() @@ -596,6 +625,28 @@ class TestOllamaConfig: assert result.choices[0]["finish_reason"] == "stop" +@pytest.mark.parametrize( + "response_body", + [ + pytest.param({}, id="missing-response"), + pytest.param({"response": "", "error": {"message": "upstream broke"}}, id="error"), + ], +) +def test_completion_with_invalid_success_reply_raises_api_connection_error( + respx_mock: respx.MockRouter, response_body: dict[str, object] +) -> None: + route: Final = respx_mock.post("http://ollama.test:11434/api/generate").respond(json=response_body) + + with pytest.raises(litellm.APIConnectionError): + litellm.completion( + model="ollama/llama3", + messages=[{"role": "user", "content": "hi"}], + api_base="http://ollama.test:11434", + ) + + assert route.called + + class TestOllamaTextCompletionResponseIterator: def test_every_chunk_of_one_stream_carries_the_same_response_id(self): iterator: Final = OllamaTextCompletionResponseIterator( @@ -682,7 +733,7 @@ class TestOllamaTextCompletionResponseIterator: # Updated to handle ModelResponseStream return type assert isinstance(result, ModelResponseStream) assert result.choices and result.choices[0].delta is not None - assert result.choices[0].delta.content == None + assert result.choices[0].delta.content is None assert getattr(result.choices[0].delta, "reasoning_content", None) == "" def test_chunk_parser_done_chunk(self):