scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

fix(oci/cohere): guard handle_cohere_response against missing usage

Browse source
This commit is contained in:
mateo-berri 2026-05-19 07:01:59 +00:00
parent 1a9e03662e
commit b59dc01159
No known key found for this signature in database
2 changed files with 24 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/llms/oci/chat/cohere.py
View file

@ -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

16
tests/test_litellm/llms/oci/test_oci_coverage_boost.py
View file

@ -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
# ===========================================================================