diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 9b5630d81ad..b65e2c498a2 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -13,7 +13,7 @@ https://platform.openai.com/docs/api-reference/batch import asyncio import contextvars import os -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import Any, Final, Literal, cast @@ -104,6 +104,24 @@ def _resolve_timeout( return float(timeout) +def _caller_timeout( + optional_params: GenericLiteLLMParams, + resolved_timeout: float | httpx.Timeout, +) -> float | httpx.Timeout | None: + """ + ``resolved_timeout`` when the caller asked for a timeout, else None. + + The resolution above falls back to 600s when nothing was set. Sending that on every + request would replace the HTTP client's own default (a 5s connect timeout, and an + explicitly configured ``litellm.request_timeout``), so None keeps that default. + """ + # ``request_timeout`` is the legacy spelling; GenericLiteLLMParams keeps it as an extra field + extra_params: Final[Mapping[str, object]] = optional_params.model_extra or {} + if optional_params.timeout or extra_params.get("request_timeout"): + return resolved_timeout + return None + + @client async def acreate_batch( completion_window: Literal["24h"], @@ -538,7 +556,7 @@ def _handle_retrieve_batch_providers_without_provider_config( batch_id=batch_id, api_base=api_base, api_key=api_key, - timeout=timeout, + timeout=_caller_timeout(optional_params, timeout), max_retries=optional_params.max_retries, litellm_params=batch_params, ) diff --git a/litellm/llms/anthropic/batches/handler.py b/litellm/llms/anthropic/batches/handler.py index a9d9f930b9e..08e4415903b 100644 --- a/litellm/llms/anthropic/batches/handler.py +++ b/litellm/llms/anthropic/batches/handler.py @@ -39,7 +39,7 @@ class AnthropicBatchesHandler: batch_id: str, api_base: str | None, api_key: str | None, - timeout: float | httpx.Timeout, + timeout: float | httpx.Timeout | None, max_retries: int | None, logging_obj: LiteLLMLoggingObj | None = None, litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment @@ -111,7 +111,7 @@ class AnthropicBatchesHandler: ) # Make the request async_client: Final = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC) - response: Final = await async_client.get(url=retrieve_url, headers=headers) + response: Final = await async_client.get(url=retrieve_url, headers=headers, timeout=timeout) response.raise_for_status() # Transform response to LiteLLM format @@ -128,7 +128,7 @@ class AnthropicBatchesHandler: batch_id: str, api_base: str | None, api_key: str | None, - timeout: float | httpx.Timeout, + timeout: float | httpx.Timeout | None, max_retries: int | None, logging_obj: LiteLLMLoggingObj | None = None, litellm_params: dict | None = None, # mutable-ok: handed straight to validate_environment diff --git a/tests/unit/llms/anthropic/batches/test_handler.py b/tests/unit/llms/anthropic/batches/test_handler.py index 9b9bddae678..688b10e0993 100644 --- a/tests/unit/llms/anthropic/batches/test_handler.py +++ b/tests/unit/llms/anthropic/batches/test_handler.py @@ -15,12 +15,14 @@ asyncio.run) is exercised directly. import asyncio import threading +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest - +import litellm +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler from litellm.types.utils import LiteLLMBatch @@ -354,3 +356,71 @@ async def test_aretrieve_batch_wif_exchange_does_not_block_event_loop(handler, p assert poster.thread_ids[0] != threading.get_ident() sent_headers = fake_client.get.call_args.kwargs["headers"] assert sent_headers["authorization"] == "Bearer sk-ant-oat01-batches-seam" + + +ANTHROPIC_BATCH_URL = "https://api.anthropic.com/v1/messages/batches/msgbatch_abc" +SEVEN_SECONDS = {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0} +# The cached client's own default: 600s, with a 5s connect timeout +CLIENT_DEFAULT = {"connect": 5.0, "read": 600.0, "write": 600.0, "pool": 600.0} + +TIMEOUT_CASES = [ + pytest.param({"timeout": 7.0}, None, SEVEN_SECONDS, id="callers-timeout"), + pytest.param({"request_timeout": 7.0}, None, SEVEN_SECONDS, id="legacy-request-timeout"), + pytest.param({}, None, CLIENT_DEFAULT, id="no-timeout-keeps-client-default"), + pytest.param( + {}, 30.0, {"connect": 5.0, "read": 30.0, "write": 30.0, "pool": 30.0}, id="no-timeout-uses-configured-timeout" + ), +] + + +@pytest.fixture +def anthropic_batch_route(respx_mock, monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + # In progress, so the success logger has no results file to download. + in_progress: Final = { + "id": "msgbatch_abc", + "processing_status": "in_progress", + "created_at": "2024-09-24T10:00:00Z", + "request_counts": {"processing": 2}, + } + return respx_mock.get(ANTHROPIC_BATCH_URL).mock(return_value=httpx.Response(200, json=in_progress)) + + +def _configure_request_timeout(monkeypatch, configured_request_timeout): + if configured_request_timeout is None: + # No global timeout configured, so a REQUEST_TIMEOUT in the environment cannot move the default + monkeypatch.setattr(litellm, "request_timeout", DEFAULT_REQUEST_TIMEOUT_SECONDS) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) + else: + monkeypatch.setattr(litellm, "request_timeout", configured_request_timeout) + + +@pytest.mark.parametrize("kwargs, configured_request_timeout, expected_timeout", TIMEOUT_CASES) +@pytest.mark.respx(assert_all_called=True) +def test_retrieve_batch_sends_the_callers_timeout_or_keeps_the_client_default( + anthropic_batch_route, monkeypatch, kwargs, configured_request_timeout, expected_timeout +): + _configure_request_timeout(monkeypatch, configured_request_timeout) + + batch = litellm.retrieve_batch( + batch_id="msgbatch_abc", custom_llm_provider="anthropic", api_key="sk-ant-test", **kwargs + ) + + assert (batch.id, batch.status) == ("msgbatch_abc", "in_progress") + assert anthropic_batch_route.calls[0].request.extensions["timeout"] == expected_timeout + + +@pytest.mark.parametrize("kwargs, configured_request_timeout, expected_timeout", TIMEOUT_CASES) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=True) +async def test_aretrieve_batch_sends_the_callers_timeout_or_keeps_the_client_default( + anthropic_batch_route, monkeypatch, kwargs, configured_request_timeout, expected_timeout +): + _configure_request_timeout(monkeypatch, configured_request_timeout) + + batch = await litellm.aretrieve_batch( + batch_id="msgbatch_abc", custom_llm_provider="anthropic", api_key="sk-ant-test", **kwargs + ) + + assert (batch.id, batch.status) == ("msgbatch_abc", "in_progress") + assert anthropic_batch_route.calls[0].request.extensions["timeout"] == expected_timeout