Merge fed96c1b50 into e16d79d635

.github/workflows/test-unit.yml

@@ -244,6 +244,7 @@ jobs:
               tests/unit/llms/clarifai
               tests/unit/llms/claude_code
               tests/unit/llms/cloudflare
+              tests/unit/llms/codestral
               tests/unit/llms/codex
               tests/unit/llms/cohere
               tests/unit/llms/cometapi

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}