From 63c0e1b9b0da68e552f86a466c3f1db21c30e2dd Mon Sep 17 00:00:00 2001 From: mountain Date: Thu, 25 Jun 2026 16:22:39 +0800 Subject: [PATCH] feat(config): pass through LiteLLM settings via a `litellm:` config block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an optional `litellm:` mapping in .openkb/config.yaml as the single place to tune LiteLLM. Keys are forwarded to LiteLLM: `timeout` and `extra_headers` apply per request (the existing per-call mechanism, covering both the compiler and the agents-SDK paths), and the rest are set as litellm module globals (drop_params, num_retries, ssl_verify, ...). Resolves the Ollama UnsupportedParamsError in #137: `litellm: {drop_params: true}` lets LiteLLM drop params a provider rejects (e.g. parallel_tool_calls on Ollama). - config.resolve_litellm_settings(): validate the value is a mapping, drop non-string keys, pass values through verbatim (the user owns them). - cli._setup_llm_key(): the `litellm:` block is canonical — when it specifies `timeout` / `extra_headers` they route to the per-call stashes and replace the legacy top-level keys (an empty `extra_headers: {}` clears, not reverts); the rest go to cli._apply_litellm_settings(), which setattrs each onto litellm. Guards: skip+warn an unknown key, and refuse to overwrite a litellm function. Globals are applied process-wide and not reset. - Back-compat: the legacy top-level `timeout:` / `extra_headers:` keys still work (now undocumented; the `litellm:` block is the documented surface). - warnings go through the logger (stderr), not click.echo, so a typo can't corrupt piped stdout (e.g. `openkb query > out.txt`). - docs: config.yaml.example consolidated under `litellm:`; README drops the now-redundant top-level extra_headers section. Closes #137 --- README.md | 10 +- config.yaml.example | 19 ++- openkb/cli.py | 42 ++++- openkb/config.py | 28 ++++ tests/test_config.py | 41 +++++ tests/test_llm_config_passthrough.py | 220 +++++++++++++++++++++++++++ 6 files changed, 340 insertions(+), 20 deletions(-) create mode 100644 tests/test_llm_config_passthrough.py diff --git a/README.md b/README.md index 093a661a..106b2b8b 100644 --- a/README.md +++ b/README.md @@ -369,19 +369,11 @@ Model names use `provider/model` LiteLLM [format](https://docs.litellm.ai/docs/p | Gemini | `gemini/gemini-3.1-pro-preview` |
-Advanced options (entity_types, extra_headers, OAuth): +Advanced options (entity_types, OAuth):
`entity_types` (optional): a YAML list overriding the entity-type vocabulary used for entity pages; omit it to use the default `person`, `organization`, `place`, `product`, `work`, `event`, `other`. -`extra_headers` (optional): a YAML mapping of extra HTTP headers sent with every LLM request (forwarded to LiteLLM's `extra_headers`). Useful for providers that expect custom headers, e.g. GitHub Copilot IDE-auth headers: - -```yaml -extra_headers: - Editor-Version: vscode/1.95.0 - Copilot-Integration-Id: vscode-chat -``` - Subscription-based providers that authenticate via OAuth device flow (e.g. `chatgpt/*`, `github_copilot/*`) need no API key; OpenKB skips the missing-key warning for them.
diff --git a/config.yaml.example b/config.yaml.example index 7efb5dd1..eebf3bf6 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -2,13 +2,6 @@ model: gpt-5.4 # LLM model (any LiteLLM-supported provider) language: en # Wiki output language pageindex_threshold: 20 # PDF pages threshold for PageIndex -# Optional: extra HTTP headers sent with every LLM request (forwarded to -# LiteLLM's extra_headers). Some providers need these — e.g. GitHub Copilot -# IDE-auth headers on older litellm versions: -# extra_headers: -# Editor-Version: vscode/1.95.0 -# Copilot-Integration-Id: vscode-chat - # Optional: override the entity-type vocabulary used for entity pages. # Omit this key to use the default 7 types # (person, organization, place, product, work, event, other). @@ -18,6 +11,12 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex # - dataset # - model -# Optional: per-request LLM timeout in seconds, forwarded to LiteLLM. -# Defaults to LiteLLM's 600s; raise it for slow local backends (e.g. Ollama). -# timeout: 1200 +# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout` 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 +# 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/openkb/cli.py b/openkb/cli.py index a34a0b95..0c25940e 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -44,6 +44,7 @@ def filter(self, record: logging.LogRecord) -> bool: from openkb.config import ( DEFAULT_CONFIG, load_config, save_config, load_global_config, register_kb, resolve_extra_headers, set_extra_headers, resolve_timeout, set_timeout, + resolve_litellm_settings, ) from openkb.converter import _registry_path, convert_document from openkb.locks import atomic_write_json, atomic_write_text, kb_ingest_lock, kb_read_lock @@ -56,6 +57,8 @@ def filter(self, record: logging.LogRecord) -> bool: load_dotenv() # load from cwd (covers running inside the KB dir) +logger = logging.getLogger(__name__) + _KNOWN_PROVIDER_KEYS = ( "OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY", @@ -83,6 +86,31 @@ def _extract_provider(model: str) -> str | None: return "openai" +def _apply_litellm_settings(settings: dict) -> None: + """Set each ``litellm:`` key verbatim onto the litellm module (process-wide + globals, so they reach every LiteLLM call). Skips with a warning a key the + installed litellm doesn't define, or one that is a litellm function (e.g. + ``completion``) since overwriting it would break later calls. Applied, never + reset — the values persist for the life of the process. + """ + for key, value in settings.items(): + if not hasattr(litellm, key): + logger.warning( + "config: LiteLLM has no setting %r — ignoring it " + "(check the spelling or your installed litellm version).", + key, + ) + continue + if callable(getattr(litellm, key)): + logger.warning( + "config: 'litellm.%s' is a LiteLLM function, not a setting — " + "refusing to overwrite it from the litellm: config block.", + key, + ) + continue + setattr(litellm, key, value) + + def _setup_llm_key(kb_dir: Path | None = None) -> None: """Set LiteLLM API key from LLM_API_KEY env var if present. @@ -113,6 +141,7 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: provider: str | None = None extra_headers: dict[str, str] = {} timeout: float | None = None + litellm_settings: dict = {} if kb_dir is not None: config_path = kb_dir / ".openkb" / "config.yaml" if config_path.exists(): @@ -121,8 +150,20 @@ 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) + 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. + 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")} + ) set_extra_headers(extra_headers) set_timeout(timeout) + _apply_litellm_settings(litellm_settings) if not api_key: # Check if any provider key is already set. OAuth-based providers @@ -304,7 +345,6 @@ def _add_single_file_locked(file_path: Path, kb_dir: Path) -> Literal["added", " from openkb.agent.compiler import compile_long_doc, compile_short_doc from openkb.state import HashRegistry - logger = logging.getLogger(__name__) openkb_dir = kb_dir / ".openkb" config = load_config(openkb_dir / "config.yaml") _setup_llm_key(kb_dir) diff --git a/openkb/config.py b/openkb/config.py index d5489688..b4cf8117 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -172,6 +172,34 @@ def resolve_timeout(config: dict) -> float | None: return value +def resolve_litellm_settings(config: dict) -> dict[str, Any]: + """Resolve the optional ``litellm:`` mapping of LiteLLM module settings. + + Values are forwarded verbatim (the user owns them); only the container shape + is enforced — returns ``{}`` if absent or not a mapping, and drops non-string + keys. ``cli._apply_litellm_settings`` applies them. + """ + raw = config.get("litellm") + if raw is None: + return {} + if not isinstance(raw, dict): + logger.warning( + "config: 'litellm' must be a mapping of LiteLLM settings, got %s — " + "ignoring it.", + type(raw).__name__, + ) + return {} + settings: dict[str, Any] = {} + for key, value in raw.items(): + if not isinstance(key, str): + logger.warning( + "config: skipping 'litellm' entry with non-string key %r.", key + ) + continue + settings[key] = value + return settings + + # Process-wide extra headers for LLM requests, resolved from the active KB's # config by the CLI entry points (cli._setup_llm_key). LLM call sites read it # via get_extra_headers() so the value doesn't have to be threaded through diff --git a/tests/test_config.py b/tests/test_config.py index 3fd8edc6..5dd870f1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,9 +1,12 @@ +import logging + from openkb.config import ( DEFAULT_CONFIG, get_extra_headers, get_timeout, load_config, resolve_extra_headers, + resolve_litellm_settings, resolve_timeout, save_config, set_extra_headers, @@ -150,3 +153,41 @@ def test_timeout_stash_roundtrip_and_reset(): assert get_timeout() == 1200.0 set_timeout(None) assert get_timeout() is None + + +def test_resolve_litellm_settings_absent_returns_empty(): + assert resolve_litellm_settings({}) == {} + + +def test_resolve_litellm_settings_passes_mapping_through_verbatim(): + # Values are forwarded as-is — no validation or coercion. + config = {"litellm": {"drop_params": True, "num_retries": 3, "ssl_verify": False}} + assert resolve_litellm_settings(config) == { + "drop_params": True, + "num_retries": 3, + "ssl_verify": False, + } + + +def test_resolve_litellm_settings_non_mapping_ignored(): + assert resolve_litellm_settings({"litellm": ["drop_params"]}) == {} + assert resolve_litellm_settings({"litellm": "drop_params=true"}) == {} + assert resolve_litellm_settings({"litellm": True}) == {} + + +def test_resolve_litellm_settings_drops_non_string_keys(): + assert resolve_litellm_settings({"litellm": {5: "x", "drop_params": True}}) == { + "drop_params": True + } + + +def test_resolve_litellm_settings_warns_on_non_mapping(caplog): + with caplog.at_level(logging.WARNING, logger="openkb.config"): + assert resolve_litellm_settings({"litellm": ["drop_params"]}) == {} + assert "must be a mapping" in caplog.text + + +def test_resolve_litellm_settings_warns_on_non_string_key(caplog): + with caplog.at_level(logging.WARNING, logger="openkb.config"): + resolve_litellm_settings({"litellm": {5: "x", "drop_params": True}}) + assert "non-string key" in caplog.text diff --git a/tests/test_llm_config_passthrough.py b/tests/test_llm_config_passthrough.py new file mode 100644 index 00000000..aa500bf2 --- /dev/null +++ b/tests/test_llm_config_passthrough.py @@ -0,0 +1,220 @@ +"""User-provided ``litellm:`` settings are applied verbatim onto the litellm module. + +``litellm:`` in config.yaml is a free-form passthrough of LiteLLM module-level +globals (``drop_params``, ``modify_params``, ``ssl_verify``, ...). +``cli._apply_litellm_settings`` assigns each key onto the ``litellm`` module so +it takes effect process-wide for every call routed through LiteLLM. It warns, +without blocking, for a key the installed LiteLLM doesn't expose (typo / version +mismatch) and refuses to overwrite a LiteLLM *function*. Settings are sticky: +applied, never reset — see ``test_apply_is_sticky_not_reset``. +""" +from __future__ import annotations + +import logging + +import litellm +import pytest + +from openkb.cli import _KNOWN_PROVIDER_KEYS, _apply_litellm_settings, _setup_llm_key + + +@pytest.fixture(autouse=True) +def _restore_litellm_globals(): + """Snapshot and restore the litellm globals these tests mutate. + + setattr on the litellm module is process-wide, so without this an applied + value would leak into every later test sharing the interpreter. ``api_key`` + is included because the end-to-end test exercises full ``_setup_llm_key``. + """ + keys = ("drop_params", "modify_params", "ssl_verify", "api_key") + saved = {k: getattr(litellm, k) for k in keys} + try: + yield + finally: + for k, v in saved.items(): + setattr(litellm, k, v) + + +def test_apply_sets_known_global_verbatim(): + litellm.drop_params = False + _apply_litellm_settings({"drop_params": True}) + assert litellm.drop_params is True + + +def test_apply_forwards_values_as_is_no_coercion(): + _apply_litellm_settings({"ssl_verify": False, "modify_params": True}) + assert litellm.ssl_verify is False + assert litellm.modify_params is True + + +def test_apply_skips_unknown_key_with_warning(caplog): + bogus = "definitely_not_a_litellm_setting_xyz" + assert not hasattr(litellm, bogus) + with caplog.at_level(logging.WARNING, logger="openkb.cli"): + _apply_litellm_settings({bogus: 123}) + # Not silently created as a dead attribute… + assert not hasattr(litellm, bogus) + # …and the user is told (on the logger, like the sibling resolvers). + assert bogus in caplog.text + assert "ignoring it" in caplog.text + + +def test_apply_refuses_to_overwrite_callable(caplog): + """A key naming a LiteLLM *function* (hasattr is True) must NOT be clobbered + — overwriting litellm.completion with a scalar would brick every later call. + """ + assert callable(litellm.completion) + with caplog.at_level(logging.WARNING, logger="openkb.cli"): + _apply_litellm_settings({"completion": 5}) + assert callable(litellm.completion) # untouched + assert "completion" in caplog.text + assert "function" in caplog.text + + +def test_apply_applies_known_even_when_another_key_is_unknown(): + litellm.drop_params = False + _apply_litellm_settings({"nope_not_real_xyz": 1, "drop_params": True}) + assert litellm.drop_params is True + + +def test_apply_empty_is_noop(): + litellm.drop_params = False + _apply_litellm_settings({}) + assert litellm.drop_params is False + + +def test_apply_is_sticky_not_reset(): + """Documented contract: settings are applied, never reset. Applying {} after + a real setting leaves the earlier value in place (it does NOT revert to the + LiteLLM default) — unlike timeout/extra_headers. Pins the intentional + process-wide stickiness so a future 'reset on empty' change is caught. + """ + litellm.drop_params = False + _apply_litellm_settings({"drop_params": True}) + _apply_litellm_settings({}) # a later config without a litellm: block + assert litellm.drop_params is True # stays set, not reset to default + + +def test_setup_llm_key_applies_litellm_block_from_config(tmp_path, monkeypatch): + """End-to-end: a ``litellm:`` block in config.yaml lands on the module the + next time any command runs _setup_llm_key. + + Env is cleared so full ``_setup_llm_key`` doesn't set litellm.api_key / + provider env vars from a key in the dev environment — which would leak + process-wide state into other tests. + """ + monkeypatch.delenv("LLM_API_KEY", raising=False) + for key in _KNOWN_PROVIDER_KEYS: + monkeypatch.delenv(key, raising=False) + + openkb_dir = tmp_path / ".openkb" + openkb_dir.mkdir(parents=True) + (openkb_dir / "config.yaml").write_text( + "model: gpt-4o-mini\nlitellm:\n drop_params: true\n", encoding="utf-8" + ) + litellm.drop_params = False + _setup_llm_key(tmp_path) + assert litellm.drop_params is True + + +def _write_kb_config(tmp_path, body: str): + openkb_dir = tmp_path / ".openkb" + openkb_dir.mkdir(parents=True) + (openkb_dir / "config.yaml").write_text(body, encoding="utf-8") + + +def _isolate_env(monkeypatch): + monkeypatch.delenv("LLM_API_KEY", raising=False) + for key in _KNOWN_PROVIDER_KEYS: + monkeypatch.delenv(key, raising=False) + + +def test_litellm_block_routes_timeout_and_extra_headers_per_call(tmp_path, monkeypatch): + """`timeout` / `extra_headers` inside the litellm: block route to the + per-call stashes (not litellm module globals); the rest stay globals. + """ + from openkb.config import get_extra_headers, get_timeout + + _isolate_env(monkeypatch) + _write_kb_config( + tmp_path, + "model: gpt-4o-mini\n" + "litellm:\n" + " timeout: 1200\n" + " drop_params: true\n" + " extra_headers:\n" + " Editor-Version: vscode/1.95.0\n", + ) + litellm.drop_params = False + _setup_llm_key(tmp_path) + assert get_timeout() == 1200.0 + assert get_extra_headers() == {"Editor-Version": "vscode/1.95.0"} + assert litellm.drop_params is True + # timeout was routed per-call, NOT setattr'd onto the module: litellm.timeout + # is still its function (a global 1200.0 would have replaced it). + assert callable(litellm.timeout) + + +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 + + _isolate_env(monkeypatch) + _write_kb_config( + tmp_path, + "model: gpt-4o-mini\n" + "timeout: 30\n" # legacy top-level + "litellm:\n" + " timeout: 1200\n", # canonical — wins + ) + _setup_llm_key(tmp_path) + assert get_timeout() == 1200.0 + + +def test_legacy_toplevel_timeout_still_works(tmp_path, monkeypatch): + """Back-compat: a top-level `timeout:` (no litellm: block) is still honored.""" + from openkb.config import get_timeout + + _isolate_env(monkeypatch) + _write_kb_config(tmp_path, "model: gpt-4o-mini\ntimeout: 900\n") + _setup_llm_key(tmp_path) + assert get_timeout() == 900.0 + + +def test_litellm_block_extra_headers_win_over_legacy_toplevel(tmp_path, monkeypatch): + """Symmetric with the timeout precedence test: a litellm: block extra_headers + replaces the legacy top-level extra_headers. + """ + from openkb.config import get_extra_headers + + _isolate_env(monkeypatch) + _write_kb_config( + tmp_path, + "model: gpt-4o-mini\n" + "extra_headers:\n" + " X-Top: toplevel\n" + "litellm:\n" + " extra_headers:\n" + " X-Block: blockval\n", + ) + _setup_llm_key(tmp_path) + assert get_extra_headers() == {"X-Block": "blockval"} + + +def test_litellm_block_empty_extra_headers_clears_legacy(tmp_path, monkeypatch): + """Regression: an explicit empty `litellm: {extra_headers: {}}` CLEARS the + legacy top-level headers, rather than silently reverting to them. + """ + from openkb.config import get_extra_headers + + _isolate_env(monkeypatch) + _write_kb_config( + tmp_path, + "model: gpt-4o-mini\n" + "extra_headers:\n" + " X-Top: toplevel\n" + "litellm:\n" + " extra_headers: {}\n", + ) + _setup_llm_key(tmp_path) + assert get_extra_headers() == {}