Merge afeff77329 into e16d79d635

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:

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":

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}]}
 

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):