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

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

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}