From 438522c1d3b399bc022dbb4b575878318880eae0 Mon Sep 17 00:00:00 2001 From: Abhishek Acharya Date: Mon, 19 Oct 2026 22:46:25 +1100 Subject: [PATCH] fix(codestral): honor timeout on FIM completions and fix sync streaming The sync, sync streaming and async streaming paths dropped the caller's timeout and used the module level client default (6000s). Only async non streaming passed it through. Sync streaming was also broken: it never sent "stream": true, so Codestral answered with a plain JSON body, and it labeled its stream wrapper "codestral" instead of "text-completion-codestral", so chunks skipped the Codestral parser and the stream raised AttributeError. --- litellm/llms/codestral/completion/handler.py | 9 ++- tests/unit/llms/codestral/__init__.py | 0 .../llms/codestral/completion/__init__.py | 0 .../test_codestral_completion_handler.py | 61 +++++++++++++++++++ 4 files changed, 68 insertions(+), 2 deletions(-) create mode 100644 tests/unit/llms/codestral/__init__.py create mode 100644 tests/unit/llms/codestral/completion/__init__.py create mode 100644 tests/unit/llms/codestral/completion/test_codestral_completion_handler.py diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 08799c3a0a5..0b91c66218f 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -103,8 +103,9 @@ async def make_call( model: str, messages: list, logging_obj, + timeout: float | httpx.Timeout | None = None, ): - response: Final = await client.post(api_base, headers=headers, data=data, stream=True) + response: Final = await client.post(api_base, headers=headers, data=data, stream=True, timeout=timeout) if response.status_code != 200: raise TextCompletionCodestralError(status_code=response.status_code, message=response.text) @@ -336,16 +337,18 @@ class CodestralTextCompletion: ### SYNC STREAMING if stream is True: + data["stream"] = True response = litellm.module_level_client.post( completion_url, headers=headers, data=json.dumps(data), stream=stream, + timeout=timeout, ) _response: Final = CustomStreamWrapper( response.iter_lines(), model, - custom_llm_provider="codestral", + custom_llm_provider="text-completion-codestral", logging_obj=logging_obj, ) return _response @@ -355,6 +358,7 @@ class CodestralTextCompletion: url=completion_url, headers=headers, data=json.dumps(data), + timeout=timeout, ) return self.process_text_completion_response( model=model, @@ -448,6 +452,7 @@ class CodestralTextCompletion: model=model, messages=messages, logging_obj=logging_obj, + timeout=timeout, ), model=model, custom_llm_provider="text-completion-codestral", diff --git a/tests/unit/llms/codestral/__init__.py b/tests/unit/llms/codestral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/__init__.py b/tests/unit/llms/codestral/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py new file mode 100644 index 00000000000..c729d733311 --- /dev/null +++ b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py @@ -0,0 +1,61 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm + +CODESTRAL_FIM_URL: Final = "https://codestral.mistral.ai/v1/fim/completions" + +COMPLETION_BODY: Final = { + "id": "cmpl-1", + "object": "chat.completion", + "created": 1, + "model": "codestral-latest", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "return 1"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}, +} + +STREAM_CHUNK: Final = { + "id": "cmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "codestral-latest", + "choices": [{"index": 0, "delta": {"content": "return 1"}, "finish_reason": "stop"}], +} + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_codestral_text_completion_sends_the_callers_timeout( + respx_mock: respx.Router, monkeypatch: pytest.MonkeyPatch, sync_mode: bool, stream: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + response: Final = ( + httpx.Response(200, content=f"data: {json.dumps(STREAM_CHUNK)}\n\ndata: [DONE]\n\n") + if stream + else httpx.Response(200, json=COMPLETION_BODY) + ) + route: Final = respx_mock.post(CODESTRAL_FIM_URL).mock(return_value=response) + + kwargs: Final = { + "model": "text-completion-codestral/codestral-latest", + "prompt": "def one():", + "api_key": "test-codestral-key", + "timeout": 7.0, + "stream": stream, + } + if sync_mode: + result = litellm.text_completion(**kwargs) + text = "".join(chunk.choices[0].text or "" for chunk in result) if stream else result.choices[0].text + else: + result = await litellm.atext_completion(**kwargs) + text = "".join([chunk.choices[0].text or "" async for chunk in result]) if stream else result.choices[0].text + + assert text == "return 1" + assert json.loads(route.calls[0].request.content).get("stream", False) is stream + assert route.calls[0].request.extensions["timeout"] == {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0}