diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index cf033e10b41..141a174e558 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -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( diff --git a/tests/unit/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py index d32b3bcbeca..fcb3d3f1cc0 100644 --- a/tests/unit/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/unit/llms/cohere/chat/test_cohere_transformation.py @@ -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")