Merge 6e3fc60f04 into aeede53112

litellm/llms/codestral/completion/handler.py

@@ -104,8 +104,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)
@@ -342,6 +343,7 @@ class CodestralTextCompletion:
                 headers=headers,
                 data=json.dumps({**data, "stream": True}),
                 stream=stream,
+                timeout=timeout,
             )
             _response: Final = CustomStreamWrapper(
                 iter_sse_lines(response.iter_text()),
@@ -356,6 +358,7 @@ class CodestralTextCompletion:
                 url=completion_url,
                 headers=headers,
                 data=json.dumps(data),
+                timeout=timeout,
             )
         return self.process_text_completion_response(
             model=model,
@@ -449,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

@@ -7,7 +7,17 @@ import respx
 
 import litellm
 
-_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
+CODESTRAL_FIM_URL: Final = "https://codestral.mistral.ai/v1/fim/completions"
+_LINE_SEPARATED_TEXT: Final = "before
between
after"
+
+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},
+}
 
 
 def _fim_stream_body(text: str) -> bytes:
@@ -40,7 +50,7 @@ async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta(
     sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
-    route: Final = respx_mock.post("https://codestral.mistral.ai/v1/fim/completions").mock(
+    route: Final = respx_mock.post(CODESTRAL_FIM_URL).mock(
         return_value=httpx.Response(200, content=_fim_stream_body(_LINE_SEPARATED_TEXT))
     )
 
@@ -48,3 +58,35 @@ async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta(
 
     assert texts == [_LINE_SEPARATED_TEXT, None]
     assert json.loads(route.calls.last.request.content)["stream"] is True
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_codestral_text_completion_sends_the_callers_timeout(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool, stream: bool
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    response: Final = (
+        httpx.Response(200, content=_fim_stream_body("return 1"))
+        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 route.calls[0].request.extensions["timeout"] == {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0}