diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 4b34c23f589..62cb71f2a52 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -228,11 +228,14 @@ def handle_cohere_response( ] usage_info = cohere_response.chatResponse.usage - model_response.usage = Usage( # type: ignore[attr-defined] - prompt_tokens=usage_info.promptTokens, # type: ignore[union-attr] - completion_tokens=usage_info.completionTokens, # type: ignore[union-attr] - total_tokens=usage_info.totalTokens, # type: ignore[union-attr] - ) + if usage_info is not None: + model_response.usage = Usage( # type: ignore[attr-defined] + prompt_tokens=usage_info.promptTokens, + completion_tokens=usage_info.completionTokens, + total_tokens=usage_info.totalTokens, + ) + else: + model_response.usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0) # type: ignore[attr-defined] return model_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 c2dc2d22251..015cd423e4f 100644 --- a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py +++ b/tests/test_litellm/llms/oci/test_oci_coverage_boost.py @@ -572,6 +572,22 @@ def test_handle_cohere_response_tool_call(): assert tool_calls[0]["function"]["name"] == "get_time" +def test_handle_cohere_response_missing_usage(): + resp = { + **_COHERE_RESPONSE_JSON, + "chatResponse": { + k: v + for k, v in _COHERE_RESPONSE_JSON["chatResponse"].items() + if k != "usage" + }, + } + model_response = ModelResponse() + result = handle_cohere_response(resp, _COHERE_MODEL, model_response) + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 + assert result.usage.total_tokens == 0 + + # =========================================================================== # cohere.py — handle_cohere_stream_chunk # ===========================================================================