From e00d360273698c44dba20cd8b1a6d0f56504ed5b Mon Sep 17 00:00:00 2001
From: mountain <kose2livs@gmail.com>
Date: Thu, 9 Jul 2026 16:40:15 +0800
Subject: [PATCH] fix: bound every LLM call with a per-request network timeout
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A stalled or half-open connection (e.g. a flaky local proxy that keeps
a socket ESTABLISHED but never sends data) could hang indexing forever
with no error — litellm/httpx had no timeout applied. Add a default
per-request timeout (120s, tunable via PAGEINDEX_LLM_TIMEOUT or
set_llm_params(timeout=...)) so a hung call fails fast with a clear
litellm Timeout, which the existing retry loops surface. Rides
get_llm_params(), so it flows through both llm_completion and
llm_acompletion automatically and is overridable per-index via
llm_params_scope / IndexConfig(llm_params=...).
---
 pageindex/config.py       | 37 +++++++++++++++++++++++++++++++++----
 tests/test_concurrency.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 4 deletions(-)

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