scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(cohere): map Cohere v2 finish_reason on non-streaming chat responses (#49529)

Browse source 
* fix(cohere): map Cohere v2 finish_reason on non-streaming chat responses

* fix(cohere): default a missing Cohere v2 finish_reason to stop instead of raising

A 200 body without finish_reason used to return stop and started raising KeyError after the mapping change, so fall back to COMPLETE
This commit is contained in:
jesus-berri 2026-11-02 12:42:23 -08:00 • committed by GitHub
parent a4de720144
commit 05b76878b0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 52 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/cohere/chat/v2_transformation.py
View file

@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
import litellm
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.types.llms.cohere import CohereV2ChatResponse
@ -248,6 +249,9 @@ class CohereV2ChatConfig(OpenAIGPTConfig):
prompt_tokens: Final = token_usage.get("input_tokens", 0)
completion_tokens: Final = token_usage.get("output_tokens", 0)
model_response.choices[0].finish_reason = map_finish_reason( # rebind-ok: ModelResponse is filled in place
cohere_v2_chat_response.get("finish_reason", "COMPLETE")
)
model_response.created = int(time.time())
model_response.model = model
usage: Final = Usage(

48
tests/unit/llms/cohere/chat/test_cohere_transformation.py
View file

@ -576,6 +576,54 @@ def test_cohere_v2_tool_calling(respx_mock: respx.MockRouter) -> None:
assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [
("get_weather_k9q2xw1n", "get_weather", '{"location":"New York, NY"}')
]
assert response.choices[0].finish_reason == "tool_calls"
@pytest.mark.usefixtures("_cohere_httpx_transport")
@pytest.mark.respx(assert_all_called=True)
@pytest.mark.parametrize(
("cohere_finish_reason", "expected_finish_reason"),
[("COMPLETE", "stop"), ("MAX_TOKENS", "length"), ("STOP_SEQUENCE", "stop")],
)
def test_cohere_v2_maps_finish_reason(
respx_mock: respx.MockRouter, cohere_finish_reason: str, expected_finish_reason: str
) -> None:
respx_mock.post(COHERE_V2_CHAT_URL).mock(
return_value=httpx.Response(
200,
json=_cohere_v2_response(
{"content": [{"type": "text", "text": "Hello"}]},
finish_reason=cohere_finish_reason,
),
)
)
response: Final = litellm.completion(
model="cohere_chat/v2/command-a-03-2025",
messages=[{"role": "user", "content": "Say hello"}],
api_key="cohere-test-key",
)
assert response.choices[0].finish_reason == expected_finish_reason
@pytest.mark.usefixtures("_cohere_httpx_transport")
@pytest.mark.respx(assert_all_called=True)
def test_cohere_v2_missing_finish_reason_defaults_to_stop(respx_mock: respx.MockRouter) -> None:
body_without_finish_reason: Final = {
key: value
for key, value in _cohere_v2_response({"content": [{"type": "text", "text": "Hello"}]}).items()
if key != "finish_reason"
}
respx_mock.post(COHERE_V2_CHAT_URL).mock(return_value=httpx.Response(200, json=body_without_finish_reason))
response: Final = litellm.completion(
model="cohere_chat/v2/command-a-03-2025",
messages=[{"role": "user", "content": "Say hello"}],
api_key="cohere-test-key",
)
assert (response.choices[0].message.content, response.choices[0].finish_reason) == ("Hello", "stop")
@pytest.mark.usefixtures("_cohere_httpx_transport")