Merge 154b03275e into e16d79d635

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,
         )

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

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