Skip to content

Commit ade8500

Browse files
committed
feat(config): pass through LiteLLM settings via a litellm: config block
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 — `timeout` / `extra_headers` in it route to the per-call stashes (and win over the legacy top-level keys); 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`). - tests + config.yaml.example. (README intentionally unchanged.) Closes #137
1 parent b2a4775 commit ade8500

5 files changed

Lines changed: 304 additions & 10 deletions

File tree

config.yaml.example

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,6 @@ model: gpt-5.4 # LLM model (any LiteLLM-supported provider)
22
language: en # Wiki output language
33
pageindex_threshold: 20 # PDF pages threshold for PageIndex
44

5-
# Optional: extra HTTP headers sent with every LLM request (forwarded to
6-
# LiteLLM's extra_headers). Some providers need these — e.g. GitHub Copilot
7-
# IDE-auth headers on older litellm versions:
8-
# extra_headers:
9-
# Editor-Version: vscode/1.95.0
10-
# Copilot-Integration-Id: vscode-chat
11-
125
# Optional: override the entity-type vocabulary used for entity pages.
136
# Omit this key to use the default 7 types
147
# (person, organization, place, product, work, event, other).
@@ -18,6 +11,12 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex
1811
# - dataset
1912
# - model
2013

21-
# Optional: per-request LLM timeout in seconds, forwarded to LiteLLM.
22-
# Defaults to LiteLLM's 600s; raise it for slow local backends (e.g. Ollama).
23-
# timeout: 1200
14+
# Optional: LLM / LiteLLM tuning. Keys are forwarded to LiteLLM; `timeout` and
15+
# `extra_headers` apply per request, the rest are set as litellm.<key>.
16+
# litellm:
17+
# timeout: 1200 # per-request timeout (s); raise for slow local backends (Ollama)
18+
# drop_params: true # let LiteLLM drop params a provider rejects (e.g. Ollama)
19+
# num_retries: 3
20+
# extra_headers: # extra HTTP headers some providers need (e.g. GitHub Copilot)
21+
# Editor-Version: vscode/1.95.0
22+
# Copilot-Integration-Id: vscode-chat

openkb/cli.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ def filter(self, record: logging.LogRecord) -> bool:
4444
from openkb.config import (
4545
DEFAULT_CONFIG, load_config, save_config, load_global_config, register_kb,
4646
resolve_extra_headers, set_extra_headers, resolve_timeout, set_timeout,
47+
resolve_litellm_settings,
4748
)
4849
from openkb.converter import _registry_path, convert_document
4950
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:
5657

5758
load_dotenv() # load from cwd (covers running inside the KB dir)
5859

60+
logger = logging.getLogger(__name__)
61+
5962

6063
_KNOWN_PROVIDER_KEYS = (
6164
"OPENAI_API_KEY", "ANTHROPIC_API_KEY", "GEMINI_API_KEY",
@@ -83,6 +86,31 @@ def _extract_provider(model: str) -> str | None:
8386
return "openai"
8487

8588

89+
def _apply_litellm_settings(settings: dict) -> None:
90+
"""Set each ``litellm:`` key verbatim onto the litellm module (process-wide
91+
globals, so they reach every LiteLLM call). Skips with a warning a key the
92+
installed litellm doesn't define, or one that is a litellm function (e.g.
93+
``completion``) since overwriting it would break later calls. Applied, never
94+
reset — the values persist for the life of the process.
95+
"""
96+
for key, value in settings.items():
97+
if not hasattr(litellm, key):
98+
logger.warning(
99+
"config: LiteLLM has no setting %r — ignoring it "
100+
"(check the spelling or your installed litellm version).",
101+
key,
102+
)
103+
continue
104+
if callable(getattr(litellm, key)):
105+
logger.warning(
106+
"config: 'litellm.%s' is a LiteLLM function, not a setting — "
107+
"refusing to overwrite it from the litellm: config block.",
108+
key,
109+
)
110+
continue
111+
setattr(litellm, key, value)
112+
113+
86114
def _setup_llm_key(kb_dir: Path | None = None) -> None:
87115
"""Set LiteLLM API key from LLM_API_KEY env var if present.
88116
@@ -113,6 +141,7 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
113141
provider: str | None = None
114142
extra_headers: dict[str, str] = {}
115143
timeout: float | None = None
144+
litellm_settings: dict = {}
116145
if kb_dir is not None:
117146
config_path = kb_dir / ".openkb" / "config.yaml"
118147
if config_path.exists():
@@ -121,8 +150,24 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
121150
provider = _extract_provider(str(model))
122151
extra_headers = resolve_extra_headers(config)
123152
timeout = resolve_timeout(config)
153+
litellm_settings = resolve_litellm_settings(config)
154+
# The litellm: block is the canonical place for LLM tuning. It may
155+
# also carry `timeout` / `extra_headers`; those route to the per-call
156+
# mechanisms (and win over the legacy top-level keys), while the rest
157+
# are applied as litellm module globals.
158+
if "extra_headers" in litellm_settings:
159+
extra_headers = resolve_extra_headers(
160+
{"extra_headers": litellm_settings.pop("extra_headers")}
161+
) or extra_headers
162+
if "timeout" in litellm_settings:
163+
block_timeout = resolve_timeout(
164+
{"timeout": litellm_settings.pop("timeout")}
165+
)
166+
if block_timeout is not None:
167+
timeout = block_timeout
124168
set_extra_headers(extra_headers)
125169
set_timeout(timeout)
170+
_apply_litellm_settings(litellm_settings)
126171

127172
if not api_key:
128173
# Check if any provider key is already set. OAuth-based providers

openkb/config.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,34 @@ def resolve_timeout(config: dict) -> float | None:
172172
return value
173173

174174

175+
def resolve_litellm_settings(config: dict) -> dict[str, Any]:
176+
"""Resolve the optional ``litellm:`` mapping of LiteLLM module settings.
177+
178+
Values are forwarded verbatim (the user owns them); only the container shape
179+
is enforced — returns ``{}`` if absent or not a mapping, and drops non-string
180+
keys. ``cli._apply_litellm_settings`` applies them.
181+
"""
182+
raw = config.get("litellm")
183+
if raw is None:
184+
return {}
185+
if not isinstance(raw, dict):
186+
logger.warning(
187+
"config: 'litellm' must be a mapping of LiteLLM settings, got %s — "
188+
"ignoring it.",
189+
type(raw).__name__,
190+
)
191+
return {}
192+
settings: dict[str, Any] = {}
193+
for key, value in raw.items():
194+
if not isinstance(key, str):
195+
logger.warning(
196+
"config: skipping 'litellm' entry with non-string key %r.", key
197+
)
198+
continue
199+
settings[key] = value
200+
return settings
201+
202+
175203
# Process-wide extra headers for LLM requests, resolved from the active KB's
176204
# config by the CLI entry points (cli._setup_llm_key). LLM call sites read it
177205
# via get_extra_headers() so the value doesn't have to be threaded through

tests/test_config.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
1+
import logging
2+
13
from openkb.config import (
24
DEFAULT_CONFIG,
35
get_extra_headers,
46
get_timeout,
57
load_config,
68
resolve_extra_headers,
9+
resolve_litellm_settings,
710
resolve_timeout,
811
save_config,
912
set_extra_headers,
@@ -150,3 +153,41 @@ def test_timeout_stash_roundtrip_and_reset():
150153
assert get_timeout() == 1200.0
151154
set_timeout(None)
152155
assert get_timeout() is None
156+
157+
158+
def test_resolve_litellm_settings_absent_returns_empty():
159+
assert resolve_litellm_settings({}) == {}
160+
161+
162+
def test_resolve_litellm_settings_passes_mapping_through_verbatim():
163+
# Values are forwarded as-is — no validation or coercion.
164+
config = {"litellm": {"drop_params": True, "num_retries": 3, "ssl_verify": False}}
165+
assert resolve_litellm_settings(config) == {
166+
"drop_params": True,
167+
"num_retries": 3,
168+
"ssl_verify": False,
169+
}
170+
171+
172+
def test_resolve_litellm_settings_non_mapping_ignored():
173+
assert resolve_litellm_settings({"litellm": ["drop_params"]}) == {}
174+
assert resolve_litellm_settings({"litellm": "drop_params=true"}) == {}
175+
assert resolve_litellm_settings({"litellm": True}) == {}
176+
177+
178+
def test_resolve_litellm_settings_drops_non_string_keys():
179+
assert resolve_litellm_settings({"litellm": {5: "x", "drop_params": True}}) == {
180+
"drop_params": True
181+
}
182+
183+
184+
def test_resolve_litellm_settings_warns_on_non_mapping(caplog):
185+
with caplog.at_level(logging.WARNING, logger="openkb.config"):
186+
assert resolve_litellm_settings({"litellm": ["drop_params"]}) == {}
187+
assert "must be a mapping" in caplog.text
188+
189+
190+
def test_resolve_litellm_settings_warns_on_non_string_key(caplog):
191+
with caplog.at_level(logging.WARNING, logger="openkb.config"):
192+
resolve_litellm_settings({"litellm": {5: "x", "drop_params": True}})
193+
assert "non-string key" in caplog.text
Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
"""User-provided ``litellm:`` settings are applied verbatim onto the litellm module.
2+
3+
``litellm:`` in config.yaml is a free-form passthrough of LiteLLM module-level
4+
globals (``drop_params``, ``modify_params``, ``ssl_verify``, ...).
5+
``cli._apply_litellm_settings`` assigns each key onto the ``litellm`` module so
6+
it takes effect process-wide for every call routed through LiteLLM. It warns,
7+
without blocking, for a key the installed LiteLLM doesn't expose (typo / version
8+
mismatch) and refuses to overwrite a LiteLLM *function*. Settings are sticky:
9+
applied, never reset — see ``test_apply_is_sticky_not_reset``.
10+
"""
11+
from __future__ import annotations
12+
13+
import logging
14+
15+
import litellm
16+
import pytest
17+
18+
from openkb.cli import _KNOWN_PROVIDER_KEYS, _apply_litellm_settings, _setup_llm_key
19+
20+
21+
@pytest.fixture(autouse=True)
22+
def _restore_litellm_globals():
23+
"""Snapshot and restore the litellm globals these tests mutate.
24+
25+
setattr on the litellm module is process-wide, so without this an applied
26+
value would leak into every later test sharing the interpreter. ``api_key``
27+
is included because the end-to-end test exercises full ``_setup_llm_key``.
28+
"""
29+
keys = ("drop_params", "modify_params", "ssl_verify", "api_key")
30+
saved = {k: getattr(litellm, k) for k in keys}
31+
try:
32+
yield
33+
finally:
34+
for k, v in saved.items():
35+
setattr(litellm, k, v)
36+
37+
38+
def test_apply_sets_known_global_verbatim():
39+
litellm.drop_params = False
40+
_apply_litellm_settings({"drop_params": True})
41+
assert litellm.drop_params is True
42+
43+
44+
def test_apply_forwards_values_as_is_no_coercion():
45+
_apply_litellm_settings({"ssl_verify": False, "modify_params": True})
46+
assert litellm.ssl_verify is False
47+
assert litellm.modify_params is True
48+
49+
50+
def test_apply_skips_unknown_key_with_warning(caplog):
51+
bogus = "definitely_not_a_litellm_setting_xyz"
52+
assert not hasattr(litellm, bogus)
53+
with caplog.at_level(logging.WARNING, logger="openkb.cli"):
54+
_apply_litellm_settings({bogus: 123})
55+
# Not silently created as a dead attribute…
56+
assert not hasattr(litellm, bogus)
57+
# …and the user is told (on the logger, like the sibling resolvers).
58+
assert bogus in caplog.text
59+
assert "ignoring it" in caplog.text
60+
61+
62+
def test_apply_refuses_to_overwrite_callable(caplog):
63+
"""A key naming a LiteLLM *function* (hasattr is True) must NOT be clobbered
64+
— overwriting litellm.completion with a scalar would brick every later call.
65+
"""
66+
assert callable(litellm.completion)
67+
with caplog.at_level(logging.WARNING, logger="openkb.cli"):
68+
_apply_litellm_settings({"completion": 5})
69+
assert callable(litellm.completion) # untouched
70+
assert "completion" in caplog.text
71+
assert "function" in caplog.text
72+
73+
74+
def test_apply_applies_known_even_when_another_key_is_unknown():
75+
litellm.drop_params = False
76+
_apply_litellm_settings({"nope_not_real_xyz": 1, "drop_params": True})
77+
assert litellm.drop_params is True
78+
79+
80+
def test_apply_empty_is_noop():
81+
litellm.drop_params = False
82+
_apply_litellm_settings({})
83+
assert litellm.drop_params is False
84+
85+
86+
def test_apply_is_sticky_not_reset():
87+
"""Documented contract: settings are applied, never reset. Applying {} after
88+
a real setting leaves the earlier value in place (it does NOT revert to the
89+
LiteLLM default) — unlike timeout/extra_headers. Pins the intentional
90+
process-wide stickiness so a future 'reset on empty' change is caught.
91+
"""
92+
litellm.drop_params = False
93+
_apply_litellm_settings({"drop_params": True})
94+
_apply_litellm_settings({}) # a later config without a litellm: block
95+
assert litellm.drop_params is True # stays set, not reset to default
96+
97+
98+
def test_setup_llm_key_applies_litellm_block_from_config(tmp_path, monkeypatch):
99+
"""End-to-end: a ``litellm:`` block in config.yaml lands on the module the
100+
next time any command runs _setup_llm_key.
101+
102+
Env is cleared so full ``_setup_llm_key`` doesn't set litellm.api_key /
103+
provider env vars from a key in the dev environment — which would leak
104+
process-wide state into other tests.
105+
"""
106+
monkeypatch.delenv("LLM_API_KEY", raising=False)
107+
for key in _KNOWN_PROVIDER_KEYS:
108+
monkeypatch.delenv(key, raising=False)
109+
110+
openkb_dir = tmp_path / ".openkb"
111+
openkb_dir.mkdir(parents=True)
112+
(openkb_dir / "config.yaml").write_text(
113+
"model: gpt-4o-mini\nlitellm:\n drop_params: true\n", encoding="utf-8"
114+
)
115+
litellm.drop_params = False
116+
_setup_llm_key(tmp_path)
117+
assert litellm.drop_params is True
118+
119+
120+
def _write_kb_config(tmp_path, body: str):
121+
openkb_dir = tmp_path / ".openkb"
122+
openkb_dir.mkdir(parents=True)
123+
(openkb_dir / "config.yaml").write_text(body, encoding="utf-8")
124+
125+
126+
def _isolate_env(monkeypatch):
127+
monkeypatch.delenv("LLM_API_KEY", raising=False)
128+
for key in _KNOWN_PROVIDER_KEYS:
129+
monkeypatch.delenv(key, raising=False)
130+
131+
132+
def test_litellm_block_routes_timeout_and_extra_headers_per_call(tmp_path, monkeypatch):
133+
"""`timeout` / `extra_headers` inside the litellm: block route to the
134+
per-call stashes (not litellm module globals); the rest stay globals.
135+
"""
136+
from openkb.config import get_extra_headers, get_timeout
137+
138+
_isolate_env(monkeypatch)
139+
_write_kb_config(
140+
tmp_path,
141+
"model: gpt-4o-mini\n"
142+
"litellm:\n"
143+
" timeout: 1200\n"
144+
" drop_params: true\n"
145+
" extra_headers:\n"
146+
" Editor-Version: vscode/1.95.0\n",
147+
)
148+
litellm.drop_params = False
149+
_setup_llm_key(tmp_path)
150+
assert get_timeout() == 1200.0
151+
assert get_extra_headers() == {"Editor-Version": "vscode/1.95.0"}
152+
assert litellm.drop_params is True
153+
# timeout/extra_headers must NOT have been setattr'd onto litellm as globals
154+
assert not isinstance(getattr(litellm, "timeout", None), float) or \
155+
getattr(litellm, "timeout", None) != 1200.0
156+
157+
158+
def test_litellm_block_timeout_wins_over_legacy_toplevel(tmp_path, monkeypatch):
159+
"""The litellm: block value wins over the legacy top-level key."""
160+
from openkb.config import get_timeout
161+
162+
_isolate_env(monkeypatch)
163+
_write_kb_config(
164+
tmp_path,
165+
"model: gpt-4o-mini\n"
166+
"timeout: 30\n" # legacy top-level
167+
"litellm:\n"
168+
" timeout: 1200\n", # canonical — wins
169+
)
170+
_setup_llm_key(tmp_path)
171+
assert get_timeout() == 1200.0
172+
173+
174+
def test_legacy_toplevel_timeout_still_works(tmp_path, monkeypatch):
175+
"""Back-compat: a top-level `timeout:` (no litellm: block) is still honored."""
176+
from openkb.config import get_timeout
177+
178+
_isolate_env(monkeypatch)
179+
_write_kb_config(tmp_path, "model: gpt-4o-mini\ntimeout: 900\n")
180+
_setup_llm_key(tmp_path)
181+
assert get_timeout() == 900.0

0 commit comments

Comments
 (0)