diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 7805f17c74c..7e380cf0e7f 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -10,6 +10,7 @@ import datetime import json from typing import Any, Dict, List, Optional +import httpx from pydantic import ValidationError from litellm.llms.oci.chat.generic import _synthesize_oci_tool_call_id @@ -202,9 +203,16 @@ def handle_cohere_response( json_response: dict, model: str, model_response: ModelResponse, + raw_response: httpx.Response, ) -> ModelResponse: """Parse a non-streaming Cohere OCI response into a LiteLLM ModelResponse.""" - cohere_response = CohereChatResult(**json_response) + try: + cohere_response = CohereChatResult(**json_response) + except (TypeError, ValidationError) as e: + raise OCIError( + message=f"Response cannot be casted to CohereChatResult: {str(e)}", + status_code=raw_response.status_code, + ) model_response.model = model model_response.created = int(datetime.datetime.now().timestamp()) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index a06c0ebdd3f..2b928220fb6 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -543,7 +543,7 @@ class OCIChatConfig(BaseConfig): vendor = get_vendor_from_model(model) if vendor == OCIVendors.COHERE: model_response = handle_cohere_response( - response_json, model, model_response + response_json, model, model_response, raw_response ) else: model_response = handle_generic_response( diff --git a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py b/tests/test_litellm/llms/oci/test_oci_coverage_boost.py index be3de91ba3d..c8934e26456 100644 --- a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py +++ b/tests/test_litellm/llms/oci/test_oci_coverage_boost.py @@ -532,10 +532,13 @@ _COHERE_RESPONSE_JSON = { } +_COHERE_RAW_RESPONSE = httpx.Response(200, request=httpx.Request("POST", "https://oci")) + + def test_handle_cohere_response_complete(): model_response = ModelResponse() result = handle_cohere_response( - _COHERE_RESPONSE_JSON, _COHERE_MODEL, model_response + _COHERE_RESPONSE_JSON, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE ) assert result.choices[0].finish_reason == "stop" assert result.choices[0].message["content"] == "Hello from Cohere!" @@ -551,7 +554,9 @@ def test_handle_cohere_response_max_tokens(): }, } model_response = ModelResponse() - result = handle_cohere_response(resp, _COHERE_MODEL, model_response) + result = handle_cohere_response( + resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE + ) assert result.choices[0].finish_reason == "length" @@ -565,7 +570,9 @@ def test_handle_cohere_response_tool_call(): }, } model_response = ModelResponse() - result = handle_cohere_response(resp, _COHERE_MODEL, model_response) + result = handle_cohere_response( + resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE + ) assert result.choices[0].finish_reason == "tool_calls" tool_calls = result.choices[0].message["tool_calls"] assert tool_calls is not None @@ -582,12 +589,23 @@ def test_handle_cohere_response_missing_usage(): }, } model_response = ModelResponse() - result = handle_cohere_response(resp, _COHERE_MODEL, model_response) + result = handle_cohere_response( + resp, _COHERE_MODEL, model_response, _COHERE_RAW_RESPONSE + ) assert result.usage.prompt_tokens == 0 assert result.usage.completion_tokens == 0 assert result.usage.total_tokens == 0 +def test_handle_cohere_response_malformed_raises_oci_error(): + bad_json = {"chatResponse": {"apiFormat": "COHERE"}} + raw = httpx.Response(502, request=httpx.Request("POST", "https://oci")) + model_response = ModelResponse() + with pytest.raises(OCIError) as exc_info: + handle_cohere_response(bad_json, _COHERE_MODEL, model_response, raw) + assert exc_info.value.status_code == 502 + + # =========================================================================== # cohere.py — handle_cohere_stream_chunk # ===========================================================================