mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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:
parent
d67cb4d969
commit
0e238785c8
2 changed files with 20 additions and 1 deletions
|
|
@ -250,7 +250,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig):
|
|||
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["finish_reason"]
|
||||
cohere_v2_chat_response.get("finish_reason", "COMPLETE")
|
||||
)
|
||||
model_response.created = int(time.time())
|
||||
model_response.model = model
|
||||
|
|
|
|||
|
|
@ -607,6 +607,25 @@ def test_cohere_v2_maps_finish_reason(
|
|||
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")
|
||||
@pytest.mark.respx(assert_all_called=True)
|
||||
async def test_cohere_v2_annotations(respx_mock: respx.MockRouter) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue