Merge 2ff4574d08 into 3b8ad37296

litellm/llms/fireworks_ai/chat/transformation.py

@@ -761,6 +761,18 @@ 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=getattr(raw_response, "headers", None),
+            )
+        if completion_fields.get("choices") is None:
+            raise FireworksAIException(
+                message=f"Fireworks AI response body is not a completion: {raw_response.text}",
+                status_code=raw_response.status_code,
+                headers=getattr(raw_response, "headers", None),
+            )
         response: Final = ModelResponse(**completion_fields)
 
         if response.model is not None:

litellm/llms/ollama/completion/transformation.py

@@ -322,6 +322,18 @@ 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" not in response_json and "message" not in response_json:
+            raise OllamaError(
+                status_code=raw_response.status_code,
+                message=f"Ollama response body is not a completion: {raw_response.text}",
+                headers=dict(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,
@@ -2358,3 +2361,36 @@ 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(
+    "body",
+    [
+        {},
+        {"choices": None},
+        {"error": {"message": "upstream broke"}},
+    ],
+)
+def test_transform_response_non_completion_body_raises(body):
+    # A 200 whose body is not a chat completion must raise instead of
+    # returning an empty successful completion.
+    config = FireworksAIConfig()
+
+    raw_response = MagicMock()
+    raw_response.status_code = 200
+    raw_response.text = json.dumps(body)
+    raw_response.headers = {}
+    raw_response.json.return_value = body
+
+    with pytest.raises(FireworksAIException):
+        config.transform_response(
+            model="accounts/fireworks/models/llama-v3p1-8b-instruct",
+            raw_response=raw_response,
+            model_response=ModelResponse(),
+            logging_obj=MagicMock(),
+            request_data={},
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            encoding=MagicMock(),
+        )

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,
@@ -990,3 +991,42 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str):
         )
         == "http://gateway.example.com/ollama/api/generate"
     )
+
+
+class TestOllamaTransformResponseMalformedBody:
+    @pytest.mark.parametrize(
+        "body",
+        [
+            {},
+            {"choices": None},
+            {"error": {"message": "upstream broke"}},
+        ],
+    )
+    def test_transform_response_non_completion_body_raises(self, body):
+        # A 200 whose body is not a generation response must raise instead of
+        # returning an empty successful completion.
+        config = OllamaConfig()
+
+        raw_response = MagicMock()
+        raw_response.status_code = 200
+        raw_response.text = json.dumps(body)
+        raw_response.headers = {}
+        raw_response.json.return_value = body
+
+        model_response = ModelResponse(
+            id="test_id",
+            choices=[{"message": Message(content="")}],
+        )
+
+        with pytest.raises(OllamaError):
+            config.transform_response(
+                model="llama2",
+                raw_response=raw_response,
+                model_response=model_response,
+                logging_obj=MagicMock(),
+                request_data={},
+                messages=[],
+                optional_params={},
+                litellm_params={},
+                encoding=MagicMock(),
+            )