diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 6b172e77295..ec1d8622ad0 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/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", diff --git a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py index 9ce50d8ceca..dac670eacb0 100644 --- a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py +++ b/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}