From 9a61576582af7b86d7d3cabdd9ff75dbe6a4a3fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=86=E5=A5=95=E4=B8=9E?= <01luyicheng@gmail.com> Date: Fri, 3 Jul 2026 15:43:54 +0000 Subject: [PATCH] feat(config): forward num_retries to LiteLLM per call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a num_retries per-call config plumbing that mirrors the existing timeout pattern. resolve_num_retries validates a non-negative int (rejects bools, floats, and strings — no coercion, unlike timeout) and a process-wide stash (set_num_retries/get_num_retries) is populated by _setup_llm_key from both the legacy top-level key and the litellm: block. The LLM call sites in compiler.py forward it to litellm.(a)completion via kwargs.setdefault. Default is None (opt-in via config) — zero behavior change for non-configurers. Refs #164. --- config.yaml.example | 7 ++- examples/configuration/README.md | 17 +++--- openkb/agent/compiler.py | 7 +++ openkb/cli.py | 18 ++++-- openkb/config.py | 42 +++++++++++++ tests/conftest.py | 8 ++- tests/test_config.py | 53 +++++++++++++++++ tests/test_llm_config_passthrough.py | 21 +++++++ tests/test_llm_retries.py | 88 ++++++++++++++++++++++++++++ 9 files changed, 244 insertions(+), 17 deletions(-) create mode 100644 tests/test_llm_retries.py diff --git a/config.yaml.example b/config.yaml.example index eebf3bf6..7036c453 100644 --- a/config.yaml.example +++ b/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.. +# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout`, +# `num_retries`, and `extra_headers` apply per request, the rest are set as +# litellm.. # 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 diff --git a/examples/configuration/README.md b/examples/configuration/README.md index 58cb7d4d..372ba3df 100644 --- a/examples/configuration/README.md +++ b/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.. +# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout`, +# `num_retries`, and `extra_headers` apply per request, the rest are set as +# litellm.. # 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) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index fc712b15..32f42f70 100644 --- a/openkb/agent/compiler.py +++ b/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) diff --git a/openkb/cli.py b/openkb/cli.py index f39d0b44..6980c275 100644 --- a/openkb/cli.py +++ b/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: diff --git a/openkb/config.py b/openkb/config.py index a5ee010b..3b4c77f1 100644 --- a/openkb/config.py +++ b/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. diff --git a/tests/conftest.py b/tests/conftest.py index 3d67b4b4..52e7d6a9 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_config.py b/tests/test_config.py index be473fcd..43c0a4cd 100644 --- a/tests/test_config.py +++ b/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({}) == {} diff --git a/tests/test_llm_config_passthrough.py b/tests/test_llm_config_passthrough.py index 7a86b73b..7bf8193e 100644 --- a/tests/test_llm_config_passthrough.py +++ b/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 diff --git a/tests/test_llm_retries.py b/tests/test_llm_retries.py new file mode 100644 index 00000000..9f2283c2 --- /dev/null +++ b/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