Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 1 addition & 9 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -369,19 +369,11 @@ Model names use `provider/model` LiteLLM [format](https://docs.litellm.ai/docs/p
| Gemini | `gemini/gemini-3.1-pro-preview` |

<details>
<summary><i>Advanced options (<code>entity_types</code>, <code>extra_headers</code>, OAuth):</i></summary>
<summary><i>Advanced options (<code>entity_types</code>, OAuth):</i></summary>
<br>

`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.

</details>
Expand Down
19 changes: 9 additions & 10 deletions config.yaml.example
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand All @@ -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.<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
# extra_headers: # extra HTTP headers some providers need (e.g. GitHub Copilot)
# Editor-Version: vscode/1.95.0
# Copilot-Integration-Id: vscode-chat
42 changes: 41 additions & 1 deletion openkb/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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",
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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():
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
28 changes: 28 additions & 0 deletions openkb/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
41 changes: 41 additions & 0 deletions tests/test_config.py
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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
Loading