diff --git a/pageindex/config.py b/pageindex/config.py index 1995d4b..dad0fc5 100644 --- a/pageindex/config.py +++ b/pageindex/config.py @@ -52,15 +52,44 @@ def _env_drop_params_default() -> bool: ) +# Built-in per-request network timeout (seconds) for every litellm completion. +# Bounds a single in-flight call so a stalled / half-open connection (e.g. a +# flaky proxy that keeps a socket ESTABLISHED but never sends data) fails fast +# with a litellm Timeout — caught by the retry loops in index/utils.py — instead +# of hanging indefinitely. Generous enough for legitimate slow responses on +# large prompts; tune via PAGEINDEX_LLM_TIMEOUT / set_llm_params(timeout=…). +_DEFAULT_LLM_TIMEOUT = 120 + + +def _env_llm_timeout_default(): + """Default per-request litellm timeout in seconds, from PAGEINDEX_LLM_TIMEOUT. + + A missing or non-numeric value falls back to ``_DEFAULT_LLM_TIMEOUT``. A + value <= 0 means "no timeout" (returns None -> litellm's own default), so a + caller can explicitly opt out. Read once at import. + """ + raw = os.getenv("PAGEINDEX_LLM_TIMEOUT", str(_DEFAULT_LLM_TIMEOUT)).strip() + try: + value = float(raw) + except ValueError: + return _DEFAULT_LLM_TIMEOUT + return value if value > 0 else None + + # Per-call kwargs PageIndex passes to every litellm completion. These are # PageIndex-OWNED and applied PER CALL — never written to litellm's shared module # globals, so they don't leak into other libraries sharing the litellm module. # Defaults preserve historical behavior: temperature=0 keeps structure # extraction deterministic; drop_params=True lets a provider that rejects a param -# (e.g. temperature on some local / reasoning models) succeed by dropping it. -# Override/extend via set_llm_params(); the common drop_params case also has the -# PAGEINDEX_DROP_PARAMS env shortcut. -_LLM_PARAMS: dict = {"temperature": 0, "drop_params": _env_drop_params_default()} +# (e.g. temperature on some local / reasoning models) succeed by dropping it; +# timeout bounds a single hung request (see _env_llm_timeout_default). Override/ +# extend via set_llm_params(); the common drop_params / timeout cases also have +# the PAGEINDEX_DROP_PARAMS / PAGEINDEX_LLM_TIMEOUT env shortcuts. +_LLM_PARAMS: dict = { + "temperature": 0, + "drop_params": _env_drop_params_default(), + "timeout": _env_llm_timeout_default(), +} # Per-call override, isolated per thread / async context — mirrors # _MAX_CONCURRENCY_OVERRIDE below. Without this, set_llm_params() is the only diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index bf2e2b8..391249b 100644 --- a/tests/test_concurrency.py +++ b/tests/test_concurrency.py @@ -7,6 +7,7 @@ import pytest from pageindex.config import ( IndexConfig, + _env_llm_timeout_default, _env_max_concurrency_default, get_llm_params, get_max_concurrency, @@ -164,6 +165,23 @@ def test_llm_acompletion_holds_the_shared_semaphore(monkeypatch): assert state["peak"] == 3 +def test_llm_acompletion_passes_a_timeout_to_litellm(monkeypatch): + # A per-request timeout must reach litellm so a hung / half-open connection + # fails fast instead of stalling indexing forever. It rides get_llm_params(), + # so both the default and an override flow through automatically. + seen = {} + + async def fake_acompletion(**kwargs): + seen.update(kwargs) + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))] + ) + + monkeypatch.setattr("litellm.acompletion", fake_acompletion) + asyncio.run(llm_acompletion("gpt-x", "hi")) + assert "timeout" in seen and seen["timeout"] == get_llm_params()["timeout"] + + def test_run_async_propagates_scope_into_worker_thread(): # When build_index runs inside an already-running loop, _run_async hops to a # worker thread. The max_concurrency_scope override must ride along (copied @@ -205,6 +223,17 @@ def test_env_default_parsing(monkeypatch): assert _env_max_concurrency_default() == 5 +def test_env_llm_timeout_parsing(monkeypatch): + monkeypatch.delenv("PAGEINDEX_LLM_TIMEOUT", raising=False) + assert _env_llm_timeout_default() == 120 + monkeypatch.setenv("PAGEINDEX_LLM_TIMEOUT", "45") + assert _env_llm_timeout_default() == 45 + monkeypatch.setenv("PAGEINDEX_LLM_TIMEOUT", "garbage") + assert _env_llm_timeout_default() == 120 + monkeypatch.setenv("PAGEINDEX_LLM_TIMEOUT", "0") # <=0 opts out of the timeout + assert _env_llm_timeout_default() is None + + def test_index_config_max_concurrency_field(): # Default is None → "use the global/env default"; explicit value overrides. assert IndexConfig().max_concurrency is None