Skip to content

Commit 936ed62

Browse files
committed
feat(config): pass through global LiteLLM settings via a litellm: config block
Add an optional `litellm:` mapping in .openkb/config.yaml whose keys are assigned verbatim onto the litellm module (litellm.<key>) — e.g. drop_params, num_retries, ssl_verify. This resolves the Ollama UnsupportedParamsError in #137: `litellm: {drop_params: true}` lets LiteLLM drop params a provider rejects (e.g. parallel_tool_calls on Ollama), which previously could only be done via the LITELLM_DROP_PARAMS env var. - config.resolve_litellm_settings(): validate the container is a mapping, drop non-string keys, pass values through unchanged (the user owns them — no per-value validation or coercion) - cli._apply_litellm_settings(): setattr each key onto litellm, applied in _setup_llm_key so it covers every command and every call path routed through LiteLLM, because these are module-level globals. Two guards: skip+warn an unknown key (typo / version mismatch can't fail silently), and refuse to overwrite a litellm *function* (e.g. completion) which would brick later calls. Settings are process-wide and sticky on purpose (applied, never reset) — documented in the docstring, README, and config.yaml.example. - warnings go through the logger (matching the sibling resolvers, on stderr), not click.echo, so a typo can't corrupt piped stdout (e.g. `openkb query >`) - tests (incl. callable-guard, sticky-no-reset, env-isolated e2e, caplog on the resolver warning paths) + config.yaml.example + README Closes #137
1 parent b2a4775 commit 936ed62

6 files changed

Lines changed: 267 additions & 0 deletions

File tree

README.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -384,6 +384,16 @@ extra_headers:
384384

385385
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.
386386

387+
`litellm` (optional): a YAML mapping of [global LiteLLM module settings](https://docs.litellm.ai/docs/), passed through verbatim — each key is assigned as `litellm.<key>`. Use it for anything LiteLLM exposes globally that OpenKB has no dedicated key for. Values are forwarded as-is (you own them); an unknown key, or a name that is a LiteLLM function rather than a setting, is warned about rather than applied. For example, to let LiteLLM drop parameters a provider rejects (e.g. Ollama not accepting `parallel_tool_calls`):
388+
389+
```yaml
390+
litellm:
391+
drop_params: true
392+
num_retries: 3
393+
```
394+
395+
These are LiteLLM module globals, so they reach every call routed through LiteLLM in the process. Two caveats: they are applied **process-wide and not reset** between commands — set them in `config.yaml`, don't rely on them reverting mid-run; and a library that sets the same global on its own calls wins for those calls (e.g. PageIndex forces `drop_params=True` per call during indexing, so setting `drop_params: false` won't make PageIndex's indexing strict).
396+
387397
</details>
388398

389399
### PageIndex Setup

config.yaml.example

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,16 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex
2121
# Optional: per-request LLM timeout in seconds, forwarded to LiteLLM.
2222
# Defaults to LiteLLM's 600s; raise it for slow local backends (e.g. Ollama).
2323
# timeout: 1200
24+
25+
# Optional: global LiteLLM module settings, passed through verbatim (each key
26+
# is set as litellm.<key>). Use this for anything LiteLLM exposes globally that
27+
# OpenKB has no dedicated key for. Values are forwarded as-is — you own them —
28+
# and a key that is unknown (or names a LiteLLM function, not a setting) is
29+
# warned about, not applied. These are module globals: applied process-wide and
30+
# NOT reset between commands, and a library that sets the same global on its own
31+
# calls wins for those (e.g. PageIndex forces drop_params=True per call while
32+
# indexing). Example: let LiteLLM drop params a provider rejects (e.g. Ollama
33+
# not accepting parallel_tool_calls):
34+
# litellm:
35+
# drop_params: true
36+
# num_retries: 3

openkb/cli.py

Lines changed: 51 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,51 @@ def _extract_provider(model: str) -> str | None:
8386
return "openai"
8487

8588

89+
def _apply_litellm_settings(settings: dict) -> None:
90+
"""Apply user-provided global LiteLLM module settings verbatim.
91+
92+
Each key is assigned onto the ``litellm`` module as-is (e.g.
93+
``litellm.drop_params``, ``litellm.num_retries``, ``litellm.ssl_verify``) —
94+
values are not validated or coerced: what the user puts in ``litellm:`` is
95+
what LiteLLM gets. Because these are module-level globals, one assignment
96+
covers every call path, including LiteLLM calls OpenKB doesn't route through
97+
its own compiler.
98+
99+
Two guards keep "verbatim" from becoming "footgun":
100+
101+
- An unknown key (one the installed LiteLLM doesn't define) is skipped with
102+
a warning rather than silently creating a dead attribute — so a typo or a
103+
litellm-version mismatch can't fail silently.
104+
- A key whose current value is *callable* (``litellm.completion``,
105+
``token_counter``, …) is refused: those are functions, not settings, and
106+
overwriting one with a config scalar would brick every later LLM call with
107+
a confusing ``'… ' object is not callable`` far from the config.
108+
109+
Process-wide and sticky on purpose: settings are applied, never reset. Once
110+
set they persist for the life of the process, so a key dropped from a later
111+
config (or a different KB read in the same process) does NOT revert to the
112+
LiteLLM default — unlike ``timeout``/``extra_headers`` which are re-resolved
113+
each call. This matches "the user owns these globals"; callers that need a
114+
clean slate must restart the process.
115+
"""
116+
for key, value in settings.items():
117+
if not hasattr(litellm, key):
118+
logger.warning(
119+
"config: LiteLLM has no setting %r — ignoring it "
120+
"(check the spelling or your installed litellm version).",
121+
key,
122+
)
123+
continue
124+
if callable(getattr(litellm, key)):
125+
logger.warning(
126+
"config: 'litellm.%s' is a LiteLLM function, not a setting — "
127+
"refusing to overwrite it from the litellm: config block.",
128+
key,
129+
)
130+
continue
131+
setattr(litellm, key, value)
132+
133+
86134
def _setup_llm_key(kb_dir: Path | None = None) -> None:
87135
"""Set LiteLLM API key from LLM_API_KEY env var if present.
88136
@@ -113,6 +161,7 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
113161
provider: str | None = None
114162
extra_headers: dict[str, str] = {}
115163
timeout: float | None = None
164+
litellm_settings: dict = {}
116165
if kb_dir is not None:
117166
config_path = kb_dir / ".openkb" / "config.yaml"
118167
if config_path.exists():
@@ -121,8 +170,10 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
121170
provider = _extract_provider(str(model))
122171
extra_headers = resolve_extra_headers(config)
123172
timeout = resolve_timeout(config)
173+
litellm_settings = resolve_litellm_settings(config)
124174
set_extra_headers(extra_headers)
125175
set_timeout(timeout)
176+
_apply_litellm_settings(litellm_settings)
126177

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

openkb/config.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,41 @@ 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 global LiteLLM module settings.
177+
178+
The value is a free-form mapping of LiteLLM module-level attribute names to
179+
values (e.g. ``drop_params: true``, ``num_retries: 3``, ``ssl_verify:
180+
false``). It is passed through verbatim — what the user configures is what
181+
LiteLLM gets — so values are NOT validated or coerced here; the user owns
182+
them. ``cli._apply_litellm_settings`` is what assigns each key onto the
183+
``litellm`` module and warns (without blocking) for a key the installed
184+
LiteLLM doesn't expose.
185+
186+
This function only enforces the container shape: returns ``{}`` when the key
187+
is absent or not a mapping, and drops (with a warning) any non-string key.
188+
"""
189+
raw = config.get("litellm")
190+
if raw is None:
191+
return {}
192+
if not isinstance(raw, dict):
193+
logger.warning(
194+
"config: 'litellm' must be a mapping of LiteLLM settings, got %s — "
195+
"ignoring it.",
196+
type(raw).__name__,
197+
)
198+
return {}
199+
settings: dict[str, Any] = {}
200+
for key, value in raw.items():
201+
if not isinstance(key, str):
202+
logger.warning(
203+
"config: skipping 'litellm' entry with non-string key %r.", key
204+
)
205+
continue
206+
settings[key] = value
207+
return settings
208+
209+
175210
# Process-wide extra headers for LLM requests, resolved from the active KB's
176211
# config by the CLI entry points (cli._setup_llm_key). LLM call sites read it
177212
# 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: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
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

0 commit comments

Comments
 (0)