diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index fa5794d71b3..b18442ac4dc 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/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: diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 239b4e32645..8c4e3a07e1f 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/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": 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..57df2621f52 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, @@ -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(), + ) diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index 248e745ebff..e299057c882 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, @@ -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(), + )