feat(config): forward num_retries to LiteLLM per call

config.yaml.example

@@ -11,12 +11,13 @@ pageindex_threshold: 20          # PDF pages threshold for PageIndex
 #   - dataset
 #   - model
 
-# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout` and
-# `extra_headers` apply per request, the rest are set as litellm.<key>.
+# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout`,
+# `num_retries`, and `extra_headers` apply per request, the rest are set as
+# litellm.<key>.
 # litellm:
 #   timeout: 1200          # per-request timeout (s); raise for slow local backends (Ollama)
 #   drop_params: true      # let LiteLLM drop params a provider rejects (e.g. Ollama)
-#   num_retries: 3
+#   num_retries: 3         # retries on transient failures (e.g. model overloaded)
 #   extra_headers:         # extra HTTP headers some providers need (e.g. GitHub Copilot)
 #     Editor-Version: vscode/1.95.0
 #     Copilot-Integration-Id: vscode-chat

examples/configuration/README.md

@@ -79,12 +79,13 @@ pageindex_threshold: 20          # PDF pages threshold for PageIndex
 #   - dataset
 #   - model
 
-# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout` and
-# `extra_headers` apply per request, the rest are set as litellm.<key>.
+# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout`,
+# `num_retries`, and `extra_headers` apply per request, the rest are set as
+# litellm.<key>.
 # litellm:
 #   timeout: 1200          # per-request timeout (s); raise for slow local backends (Ollama)
 #   drop_params: true      # let LiteLLM drop params a provider rejects (e.g. Ollama)
-#   num_retries: 3
+#   num_retries: 3         # retries on transient failures (e.g. model overloaded)
 #   extra_headers:         # extra HTTP headers some providers need (e.g. GitHub Copilot)
 #     Editor-Version: vscode/1.95.0
 #     Copilot-Integration-Id: vscode-chat
@@ -101,12 +102,12 @@ pageindex_threshold: 20          # PDF pages threshold for PageIndex
 ### The `litellm:` block
 
 OpenKB forwards this block to LiteLLM so you can tune anything LiteLLM supports —
-you set it, LiteLLM uses it. Two keys are special:
+you set it, LiteLLM uses it. Three keys are special:
 
-- `timeout` and `extra_headers` are applied **per request** (they're needed on
-  every call).
-- Every other key (`drop_params`, `num_retries`, `ssl_verify`, …) is set on the
-  `litellm` module as a process-wide global.
+- `timeout`, `num_retries`, and `extra_headers` are applied **per request**
+  (they're needed on every call).
+- Every other key (`drop_params`, `ssl_verify`, …) is set on the `litellm`
+  module as a process-wide global.
 
 #### Slow local runtimes (Ollama, LM Studio, llama.cpp)
 

openkb/agent/compiler.py

@@ -33,6 +33,7 @@
 from openkb.config import (
     DEFAULT_ENTITY_TYPES,
     get_extra_headers,
+    get_num_retries,
     get_timeout,
     resolve_entity_types,
 )
@@ -408,6 +409,9 @@ def _llm_call(
     timeout = get_timeout()
     if timeout is not None:
         kwargs.setdefault("timeout", timeout)
+    num_retries = get_num_retries()
+    if num_retries is not None:
+        kwargs.setdefault("num_retries", num_retries)
     logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages))
     if kwargs:
         logger.debug("LLM kwargs [%s]: %s", step_name, kwargs)
@@ -442,6 +446,9 @@ async def _llm_call_async(
     timeout = get_timeout()
     if timeout is not None:
         kwargs.setdefault("timeout", timeout)
+    num_retries = get_num_retries()
+    if num_retries is not None:
+        kwargs.setdefault("num_retries", num_retries)
     logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages))
     if kwargs:
         logger.debug("LLM kwargs [%s]: %s", step_name, kwargs)

openkb/cli.py

@@ -58,6 +58,8 @@ def filter(self, record: logging.LogRecord) -> bool:
     set_extra_headers,
     resolve_timeout,
     set_timeout,
+    resolve_num_retries,
+    set_num_retries,
     resolve_litellm_settings,
 )
 from openkb.add_coordinator import _cleanup_staging_dirs
@@ -169,11 +171,12 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
 
     api_key = os.environ.get("LLM_API_KEY", "")
 
-    # Try to resolve the active provider, extra headers, and request timeout
-    # from the KB config
+    # Try to resolve the active provider, extra headers, request timeout, and
+    # retry count from the KB config
     provider: str | None = None
     extra_headers: dict[str, str] = {}
     timeout: float | None = None
+    num_retries: int | None = None
     litellm_settings: dict = {}
     if kb_dir is not None:
         config_path = kb_dir / ".openkb" / "config.yaml"
@@ -183,17 +186,24 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
             provider = _extract_provider(str(model))
             extra_headers = resolve_extra_headers(config)
             timeout = resolve_timeout(config)
+            num_retries = resolve_num_retries(config)
             litellm_settings = resolve_litellm_settings(config)
-            # `timeout` / `extra_headers` in the block route to the per-call
-            # stashes (replacing the legacy top-level keys); the rest are globals.
+            # `timeout` / `extra_headers` / `num_retries` in the block route to
+            # the per-call stashes (replacing the legacy top-level keys); the
+            # rest are globals.
             if "extra_headers" in litellm_settings:
                 extra_headers = resolve_extra_headers(
                     {"extra_headers": litellm_settings.pop("extra_headers")}
                 )
             if "timeout" in litellm_settings:
                 timeout = resolve_timeout({"timeout": litellm_settings.pop("timeout")})
+            if "num_retries" in litellm_settings:
+                num_retries = resolve_num_retries(
+                    {"num_retries": litellm_settings.pop("num_retries")}
+                )
     set_extra_headers(extra_headers)
     set_timeout(timeout)
+    set_num_retries(num_retries)
     _apply_litellm_settings(litellm_settings)
 
     if not api_key:

openkb/config.py

@@ -176,6 +176,32 @@ def resolve_timeout(config: dict) -> float | None:
     return value
 
 
+def resolve_num_retries(config: dict) -> int | None:
+    """Resolve the optional ``num_retries:`` key to a non-negative int.
+
+    Returns ``None`` (use LiteLLM's default) when absent or invalid; rejects
+    bools, floats, and strings, warning when present but unusable. A retry
+    count is a discrete count, so no numeric coercion is applied (unlike
+    :func:`resolve_timeout`).
+    """
+    raw = config.get("num_retries")
+    if raw is None:
+        return None
+    if isinstance(raw, bool) or not isinstance(raw, int):
+        logger.warning(
+            "config: 'num_retries' must be a non-negative integer, got %s — ignoring it.",
+            type(raw).__name__,
+        )
+        return None
+    if raw < 0:
+        logger.warning(
+            "config: 'num_retries' must be a non-negative integer, got %s — ignoring it.",
+            raw,
+        )
+        return None
+    return raw
+
+
 def resolve_litellm_settings(config: dict) -> dict[str, Any]:
     """Resolve the optional ``litellm:`` mapping of LiteLLM module settings.
 
@@ -243,6 +269,22 @@ def get_timeout_extra_args() -> dict[str, float] | None:
     return {"timeout": _runtime_timeout} if _runtime_timeout is not None else None
 
 
+# Process-wide LLM request retry count, set from config by the CLI and read at
+# the call sites via get_num_retries(). None = use LiteLLM's default.
+_runtime_num_retries: int | None = None
+
+
+def set_num_retries(num_retries: int | None) -> None:
+    """Set the process-wide LLM request retry count; ``None`` clears it."""
+    global _runtime_num_retries
+    _runtime_num_retries = num_retries
+
+
+def get_num_retries() -> int | None:
+    """Return the process-wide LLM request retry count, or ``None``."""
+    return _runtime_num_retries
+
+
 def load_config(config_path: Path) -> dict[str, Any]:
     """Load YAML config from config_path, merged with DEFAULT_CONFIG.
 

tests/conftest.py

@@ -5,12 +5,16 @@
 
 @pytest.fixture(autouse=True)
 def _reset_extra_headers():
-    """Keep the process-wide LLM extra-headers / timeout stashes from leaking across tests."""
-    from openkb.config import set_extra_headers, set_timeout
+    """Keep process-wide LLM stashes from leaking across tests.
+
+    Resets extra-headers, timeout, and num_retries between tests.
+    """
+    from openkb.config import set_extra_headers, set_num_retries, set_timeout
 
     yield
     set_extra_headers({})
     set_timeout(None)
+    set_num_retries(None)
 
 
 @pytest.fixture

tests/test_config.py

@@ -3,13 +3,16 @@
 from openkb.config import (
     DEFAULT_CONFIG,
     get_extra_headers,
+    get_num_retries,
     get_timeout,
     load_config,
     resolve_extra_headers,
     resolve_litellm_settings,
+    resolve_num_retries,
     resolve_timeout,
     save_config,
     set_extra_headers,
+    set_num_retries,
     set_timeout,
 )
 
@@ -161,6 +164,56 @@ def test_timeout_stash_roundtrip_and_reset():
     assert get_timeout() is None
 
 
+# --- num_retries -------------------------------------------------------------
+
+
+def test_resolve_num_retries_absent_returns_none():
+    assert resolve_num_retries({}) is None
+
+
+def test_resolve_num_retries_int():
+    assert resolve_num_retries({"num_retries": 3}) == 3
+    assert resolve_num_retries({"num_retries": 1}) == 1
+
+
+def test_resolve_num_retries_zero_allowed():
+    # 0 is a valid retry count (no retries, but explicitly set).
+    assert resolve_num_retries({"num_retries": 0}) == 0
+
+
+def test_resolve_num_retries_rejects_negative():
+    assert resolve_num_retries({"num_retries": -1}) is None
+
+
+def test_resolve_num_retries_rejects_bool():
+    # bool is a subclass of int; True/False are not retry counts.
+    assert resolve_num_retries({"num_retries": True}) is None
+    assert resolve_num_retries({"num_retries": False}) is None
+
+
+def test_resolve_num_retries_rejects_float():
+    # A retry count is discrete; floats are rejected (no coercion).
+    assert resolve_num_retries({"num_retries": 3.5}) is None
+    assert resolve_num_retries({"num_retries": 3.0}) is None
+
+
+def test_resolve_num_retries_rejects_string():
+    # Unlike timeout, numeric strings are NOT coerced.
+    assert resolve_num_retries({"num_retries": "3"}) is None
+
+
+def test_resolve_num_retries_rejects_non_int():
+    assert resolve_num_retries({"num_retries": [3]}) is None
+    assert resolve_num_retries({"num_retries": {"a": 1}}) is None
+
+
+def test_num_retries_stash_roundtrip_and_reset():
+    set_num_retries(3)
+    assert get_num_retries() == 3
+    set_num_retries(None)
+    assert get_num_retries() is None
+
+
 def test_resolve_litellm_settings_absent_returns_empty():
     assert resolve_litellm_settings({}) == {}
 

tests/test_llm_config_passthrough.py

@@ -156,6 +156,27 @@ def test_litellm_block_routes_timeout_and_extra_headers_per_call(tmp_path, monke
     assert callable(litellm.timeout)
 
 
+def test_litellm_block_routes_num_retries_per_call(tmp_path, monkeypatch):
+    """`num_retries` inside the litellm: block routes to the per-call stash
+    (not a litellm module global); the rest stay globals.
+    """
+    from openkb.config import get_num_retries
+
+    _isolate_env(monkeypatch)
+    _write_kb_config(
+        tmp_path,
+        "model: gpt-4o-mini\nlitellm:\n  num_retries: 4\n  drop_params: true\n",
+    )
+    litellm.drop_params = False
+    before = getattr(litellm, "num_retries", None)
+    _setup_llm_key(tmp_path)
+    assert get_num_retries() == 4
+    assert litellm.drop_params is True
+    # num_retries was routed per-call, NOT setattr'd onto the module: the
+    # module-level value is unchanged (a global 4 would have replaced it).
+    assert getattr(litellm, "num_retries", None) == before
+
+
 def test_litellm_block_timeout_wins_over_legacy_toplevel(tmp_path, monkeypatch):
     """The litellm: block value wins over the legacy top-level key."""
     from openkb.config import get_timeout

tests/test_llm_retries.py

@@ -0,0 +1,88 @@
+"""The configurable LLM request retry count is forwarded to LiteLLM.
+
+`num_retries:` in config.yaml is resolved into a process-wide stash (see
+test_config.py) and read at the LiteLLM call sites in openkb.agent.compiler.
+These tests pin the call-site behavior: a configured retry count is forwarded to
+`litellm.(a)completion`, and nothing is forwarded when it is unset (so LiteLLM
+keeps applying its own default).
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from openkb.agent.compiler import _llm_call, _llm_call_async
+from openkb.config import set_num_retries
+
+
+def _fake_response():
+    choice = MagicMock()
+    choice.message.content = "ok"
+    choice.finish_reason = "stop"
+    resp = MagicMock()
+    resp.choices = [choice]
+    return resp
+
+
+def test_llm_call_forwards_configured_num_retries():
+    set_num_retries(3)
+    with patch(
+        "openkb.agent.compiler.litellm.completion", return_value=_fake_response()
+    ) as completion:
+        _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step")
+    assert completion.call_args.kwargs["num_retries"] == 3
+
+
+def test_llm_call_omits_num_retries_when_unset():
+    set_num_retries(None)
+    with patch(
+        "openkb.agent.compiler.litellm.completion", return_value=_fake_response()
+    ) as completion:
+        _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step")
+    assert "num_retries" not in completion.call_args.kwargs
+
+
+def test_llm_call_does_not_override_explicit_num_retries():
+    # An explicit per-call num_retries kwarg wins over the configured default.
+    set_num_retries(3)
+    with patch(
+        "openkb.agent.compiler.litellm.completion", return_value=_fake_response()
+    ) as completion:
+        _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step", num_retries=5)
+    assert completion.call_args.kwargs["num_retries"] == 5
+
+
+def test_llm_call_async_forwards_configured_num_retries():
+    set_num_retries(2)
+    with patch(
+        "openkb.agent.compiler.litellm.acompletion",
+        new_callable=AsyncMock,
+        return_value=_fake_response(),
+    ) as acompletion:
+        asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step"))
+    assert acompletion.call_args.kwargs["num_retries"] == 2
+
+
+def test_llm_call_async_omits_num_retries_when_unset():
+    set_num_retries(None)
+    with patch(
+        "openkb.agent.compiler.litellm.acompletion",
+        new_callable=AsyncMock,
+        return_value=_fake_response(),
+    ) as acompletion:
+        asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step"))
+    assert "num_retries" not in acompletion.call_args.kwargs
+
+
+def test_llm_call_async_does_not_override_explicit_num_retries():
+    set_num_retries(2)
+    with patch(
+        "openkb.agent.compiler.litellm.acompletion",
+        new_callable=AsyncMock,
+        return_value=_fake_response(),
+    ) as acompletion:
+        asyncio.run(
+            _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step", num_retries=5)
+        )
+    assert acompletion.call_args.kwargs["num_retries"] == 5