From fee446087c6ec096fbf0dbfa3ffbc88693790d6c Mon Sep 17 00:00:00 2001 From: mountain Date: Sat, 13 Jun 2026 09:49:10 +0800 Subject: [PATCH 1/2] fix(deps): pin all dependencies to exact versions; bump litellm to 1.87.2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit litellm 1.87.2 fixes the chatgpt/* (ChatGPT subscription) provider returning empty Responses output (BerriAI/litellm#25429) and auto-injects GitHub Copilot IDE-auth headers in the chat path (BerriAI/litellm#20113). All direct dependencies are now pinned exactly as a supply-chain precaution (cf. the litellm package-poisoning incident) — the README already claimed litellm was pinned to a safe version, but pyproject carried no constraint at all. Versions match what uv.lock already resolved, so installed behavior is unchanged apart from litellm 1.85.0 -> 1.87.2. Refs VectifyAI/OpenKB#94 --- pyproject.toml | 30 ++++++++++++++++++------------ uv.lock | 32 ++++++++++++++++---------------- 2 files changed, 34 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 026dea23b..df77078ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,19 +25,25 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] keywords = ["ai", "rag", "retrieval", "knowledge-base", "llm", "pageindex", "agents", "document"] +# All dependencies are pinned exactly (supply-chain caution — e.g. the +# litellm package-poisoning incident). Bump deliberately after vetting +# each release. +# litellm 1.87.2 fixes the chatgpt/* (ChatGPT subscription) provider +# returning empty Responses output (BerriAI/litellm#25429) and +# auto-injects GitHub Copilot IDE-auth headers. dependencies = [ "pageindex==0.3.0.dev1", - "markitdown[docx,pptx,xlsx,xls]>=0.1.5", - "trafilatura>=2.0", - "click>=8.0", - "watchdog>=3.0", - "litellm", - "openai-agents", - "pyyaml", - "python-dotenv", - "json-repair", - "prompt_toolkit>=3.0", - "rich>=13.0", + "markitdown[docx,pptx,xlsx,xls]==0.1.5", + "trafilatura==2.0.0", + "click==8.4.0", + "watchdog==6.0.0", + "litellm==1.87.2", + "openai-agents==0.17.3", + "pyyaml==6.0.3", + "python-dotenv==1.2.2", + "json-repair==0.59.10", + "prompt_toolkit==3.0.52", + "rich==15.0.0", ] [project.urls] @@ -52,7 +58,7 @@ openkb = "openkb.cli:cli" testpaths = ["tests"] [project.optional-dependencies] -dev = ["pytest", "pytest-asyncio"] +dev = ["pytest==9.0.3", "pytest-asyncio==1.3.0"] [tool.hatch.version] source = "vcs" diff --git a/uv.lock b/uv.lock index ea1d268e3..227e6b1ef 100644 --- a/uv.lock +++ b/uv.lock @@ -1148,7 +1148,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.85.0" +version = "1.87.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -1164,9 +1164,9 @@ dependencies = [ { name = "tiktoken" }, { name = "tokenizers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3b/d5/3c9b560db2ffa9e498655d0dfd74f408bc5b32ede858b5731c2a5fa4c752/litellm-1.85.0.tar.gz", hash = "sha256:babdd569809af913d08a08a7eb55df1ed3e6a3960ee365c6cef4ad031c9bc72a", size = 15344387, upload-time = "2026-05-17T01:59:15.97Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/63/a8623c440a2e2de401c2dd531c376fc74180908f3c1dd4716544b8610cd0/litellm-1.87.2.tar.gz", hash = "sha256:36b172cf8824bed6ae17fb8031a021c8bb2046014303c2220045868f1230690a", size = 15459811, upload-time = "2026-06-11T05:00:02.098Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1c/38/e6a4abb062e039d18d59538cc4e6fc370c2c10cd2bff4a2e546acb69dcb9/litellm-1.85.0-py3-none-any.whl", hash = "sha256:2bb449153610691faffd76f5b94a8c29e4b66fc5394156ebf54fd4fe92759b1a", size = 16978229, upload-time = "2026-05-17T01:59:11.902Z" }, + { url = "https://files.pythonhosted.org/packages/46/5d/aac1831ada451a3fdd27ed26548ef775a8b29751ac79ea1a105ce693d628/litellm-1.87.2-py3-none-any.whl", hash = "sha256:d2161fc0ce66f800c69fd522cb73f03683edd164fc668e00fc3f23e1f6b4e6bf", size = 17111020, upload-time = "2026-06-11T04:59:59.395Z" }, ] [[package]] @@ -1918,20 +1918,20 @@ dev = [ [package.metadata] requires-dist = [ - { name = "click", specifier = ">=8.0" }, - { name = "json-repair" }, - { name = "litellm" }, - { name = "markitdown", extras = ["docx", "pptx", "xls", "xlsx"], specifier = ">=0.1.5" }, - { name = "openai-agents" }, + { name = "click", specifier = "==8.4.0" }, + { name = "json-repair", specifier = "==0.59.10" }, + { name = "litellm", specifier = "==1.87.2" }, + { name = "markitdown", extras = ["docx", "pptx", "xls", "xlsx"], specifier = "==0.1.5" }, + { name = "openai-agents", specifier = "==0.17.3" }, { name = "pageindex", specifier = "==0.3.0.dev1" }, - { name = "prompt-toolkit", specifier = ">=3.0" }, - { name = "pytest", marker = "extra == 'dev'" }, - { name = "pytest-asyncio", marker = "extra == 'dev'" }, - { name = "python-dotenv" }, - { name = "pyyaml" }, - { name = "rich", specifier = ">=13.0" }, - { name = "trafilatura", specifier = ">=2.0" }, - { name = "watchdog", specifier = ">=3.0" }, + { name = "prompt-toolkit", specifier = "==3.0.52" }, + { name = "pytest", marker = "extra == 'dev'", specifier = "==9.0.3" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = "==1.3.0" }, + { name = "python-dotenv", specifier = "==1.2.2" }, + { name = "pyyaml", specifier = "==6.0.3" }, + { name = "rich", specifier = "==15.0.0" }, + { name = "trafilatura", specifier = "==2.0.0" }, + { name = "watchdog", specifier = "==6.0.0" }, ] provides-extras = ["dev"] From 51a3d057ae5e984f3470dca1ffbeb6009c7c2915 Mon Sep 17 00:00:00 2001 From: mountain Date: Sat, 13 Jun 2026 09:49:21 +0800 Subject: [PATCH 2/2] feat(llm): forward config extra_headers to all LLM calls; no key warning for OAuth providers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Issue #93: users on GitHub Copilot (and similar providers) need custom HTTP headers (Editor-Version etc.) on every LLM request, and adding extra_headers to config.yaml silently did nothing — the key was loaded into the config dict but never forwarded. config.resolve_extra_headers() validates/stringifies the mapping, and cli._setup_llm_key (called by every LLM-using command) stashes it via config.set_extra_headers(). The two call funnels read the stash: - compiler._llm_call/_llm_call_async pass it as litellm extra_headers (explicit per-call kwargs still win), - every agents-SDK Agent gets ModelSettings(extra_headers=...), which LitellmModel forwards to litellm.acompletion; build_chat_agent and skill_runner inherit it via base.clone(). Issue #94: OAuth/subscription providers (chatgpt/*, github_copilot/*) authenticate via device flow and need no API key, so the 'No LLM API key found' warning is skipped for them. Closes VectifyAI/OpenKB#93 Refs VectifyAI/OpenKB#94 --- README.md | 10 +++++ config.yaml.example | 7 ++++ openkb/agent/compiler.py | 8 +++- openkb/agent/linter.py | 3 ++ openkb/agent/query.py | 6 ++- openkb/cli.py | 21 +++++++++-- openkb/config.py | 60 +++++++++++++++++++++++++++++ openkb/skill/creator.py | 6 ++- openkb/skill/evaluator.py | 5 +++ tests/conftest.py | 8 ++++ tests/test_cli.py | 79 +++++++++++++++++++++++++++++++++++++++ tests/test_compiler.py | 51 +++++++++++++++++++++++++ tests/test_config.py | 59 ++++++++++++++++++++++++++++- tests/test_query.py | 19 ++++++++++ 14 files changed, 334 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 600a58152..cfffbc83f 100644 --- a/README.md +++ b/README.md @@ -273,6 +273,16 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex `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. + Model names use `provider/model` LiteLLM [format](https://docs.litellm.ai/docs/providers) (OpenAI models can omit the prefix): | Provider | Model example | diff --git a/config.yaml.example b/config.yaml.example index 9054b9773..45b4b8c36 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -2,6 +2,13 @@ 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). diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index f904f9ee9..8a4226828 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -29,7 +29,7 @@ import litellm import yaml -from openkb.config import DEFAULT_ENTITY_TYPES, resolve_entity_types +from openkb.config import DEFAULT_ENTITY_TYPES, get_extra_headers, resolve_entity_types from openkb.lint import list_existing_wiki_targets, strip_ghost_wikilinks from openkb.schema import INDEX_SEED, get_agents_md @@ -323,6 +323,9 @@ def _fmt_messages(messages: list[dict], max_content: int = 200) -> str: def _llm_call(model: str, messages: list[dict], step_name: str, **kwargs) -> str: """Single LLM call with animated progress and debug logging.""" + extra_headers = get_extra_headers() + if extra_headers: + kwargs.setdefault("extra_headers", extra_headers) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -342,6 +345,9 @@ def _llm_call(model: str, messages: list[dict], step_name: str, **kwargs) -> str async def _llm_call_async(model: str, messages: list[dict], step_name: str, **kwargs) -> str: """Async LLM call with timing output and debug logging.""" + extra_headers = get_extra_headers() + if extra_headers: + kwargs.setdefault("extra_headers", extra_headers) 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/agent/linter.py b/openkb/agent/linter.py index 7e790299f..7e9ae8e77 100644 --- a/openkb/agent/linter.py +++ b/openkb/agent/linter.py @@ -4,8 +4,10 @@ from pathlib import Path from agents import Agent, Runner, function_tool +from agents.model_settings import ModelSettings from openkb.agent.tools import list_wiki_files, read_wiki_file +from openkb.config import get_extra_headers MAX_TURNS = 50 from openkb.schema import get_agents_md @@ -79,6 +81,7 @@ def read_file(path: str) -> str: instructions=instructions, tools=[list_files, read_file], model=f"litellm/{model}", + model_settings=ModelSettings(extra_headers=get_extra_headers() or None), ) diff --git a/openkb/agent/query.py b/openkb/agent/query.py index 219277cc7..a8a97ce5d 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -6,6 +6,7 @@ from agents import Agent, Runner, function_tool from agents import ToolOutputImage, ToolOutputText +from openkb.config import get_extra_headers from openkb.agent.tools import ( get_wiki_page_content, read_wiki_file, @@ -94,7 +95,10 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: instructions=instructions, tools=[read_file, get_page_content, get_image], model=f"litellm/{model}", - model_settings=ModelSettings(parallel_tool_calls=False), + model_settings=ModelSettings( + parallel_tool_calls=False, + extra_headers=get_extra_headers() or None, + ), ) diff --git a/openkb/cli.py b/openkb/cli.py index b963fc204..a36140575 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -41,7 +41,10 @@ def filter(self, record: logging.LogRecord) -> bool: litellm.suppress_debug_info = True from dotenv import load_dotenv -from openkb.config import DEFAULT_CONFIG, load_config, save_config, load_global_config, register_kb +from openkb.config import ( + DEFAULT_CONFIG, load_config, save_config, load_global_config, register_kb, + resolve_extra_headers, set_extra_headers, +) from openkb.converter import _registry_path, convert_document from openkb.locks import atomic_write_json, atomic_write_text, kb_ingest_lock, kb_read_lock from openkb.log import append_log @@ -60,6 +63,11 @@ def filter(self, record: logging.LogRecord) -> bool: "ZHIPUAI_API_KEY", "DASHSCOPE_API_KEY", ) +# Providers that authenticate via OAuth device flow (subscription login +# handled by LiteLLM itself) — no API key env var is needed, so the +# missing-key warning would be a false alarm for them. +_OAUTH_PROVIDERS = {"chatgpt", "github_copilot"} + def _extract_provider(model: str) -> str | None: """Extract the LiteLLM provider name from a model string. @@ -100,23 +108,28 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: api_key = os.environ.get("LLM_API_KEY", "") - # Try to resolve the active provider from the KB config + # Try to resolve the active provider and extra headers from the KB config provider: str | None = None + extra_headers: dict[str, str] = {} if kb_dir is not None: config_path = kb_dir / ".openkb" / "config.yaml" if config_path.exists(): config = load_config(config_path) model = config.get("model", "") provider = _extract_provider(str(model)) + extra_headers = resolve_extra_headers(config) + set_extra_headers(extra_headers) if not api_key: - # Check if any provider key is already set + # Check if any provider key is already set. OAuth-based providers + # (ChatGPT subscription, GitHub Copilot) don't use API keys at all, + # so the warning is skipped for them. check_keys = ( (f"{provider.upper()}_API_KEY",) if provider else _KNOWN_PROVIDER_KEYS ) has_key = any(os.environ.get(k) for k in check_keys) - if not has_key: + if not has_key and provider not in _OAUTH_PROVIDERS: click.echo( "Warning: No LLM API key found. Set one of:\n" f" 1. {kb_dir / '.env' if kb_dir else '/.env'} — LLM_API_KEY=sk-...\n" diff --git a/openkb/config.py b/openkb/config.py index 223ede85e..24108a7e5 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -96,6 +96,66 @@ def resolve_entity_types(config: dict) -> list[str]: return cleaned +def resolve_extra_headers(config: dict) -> dict[str, str]: + """Resolve the optional ``extra_headers:`` config key into a str→str dict. + + Some LiteLLM providers need extra HTTP headers on every request (e.g. + GitHub Copilot's ``Editor-Version`` IDE-auth headers). Users opt in via + an ``extra_headers:`` mapping in config.yaml; the result is forwarded to + LiteLLM's ``extra_headers`` parameter on all LLM calls. + + Values are stringified (YAML may parse version-like values as numbers). + Entries with a non-string/empty key or a non-scalar value are skipped. + A non-mapping ``extra_headers`` is ignored entirely. Warnings are logged + only when the key was present but malformed. + """ + raw = config.get("extra_headers") + if raw is None: + return {} + if not isinstance(raw, dict): + logger.warning( + "config: 'extra_headers' must be a mapping of header name to " + "value, got %s — ignoring it.", + type(raw).__name__, + ) + return {} + headers: dict[str, str] = {} + for key, value in raw.items(): + if not isinstance(key, str) or not key.strip(): + logger.warning( + "config: skipping 'extra_headers' entry with non-string " + "or empty key: %r", key, + ) + continue + if value is None or not isinstance(value, (str, int, float, bool)): + logger.warning( + "config: skipping 'extra_headers' entry %r with " + "non-scalar value: %r", key, value, + ) + continue + headers[key.strip()] = str(value) + return headers + + +# 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 +# every compile/agent call chain — mirroring how the API key is applied +# globally via litellm.api_key / provider env vars. +_runtime_extra_headers: dict[str, str] = {} + + +def set_extra_headers(headers: dict[str, str]) -> None: + """Set the process-wide extra headers for LLM requests.""" + global _runtime_extra_headers + _runtime_extra_headers = dict(headers) + + +def get_extra_headers() -> dict[str, str]: + """Return a copy of the process-wide extra headers for LLM requests.""" + return dict(_runtime_extra_headers) + + def load_config(config_path: Path) -> dict[str, Any]: """Load YAML config from config_path, merged with DEFAULT_CONFIG. diff --git a/openkb/skill/creator.py b/openkb/skill/creator.py index 4158d81d8..7a9d68b1c 100644 --- a/openkb/skill/creator.py +++ b/openkb/skill/creator.py @@ -19,6 +19,7 @@ from agents import Agent, Runner, ToolOutputImage, ToolOutputText, function_tool from agents.model_settings import ModelSettings +from openkb.config import get_extra_headers from openkb.skill import skill_dir from openkb.skill.tools import ( get_skill_page_content as _get_page_content_impl, @@ -150,7 +151,10 @@ def done(summary: str) -> str: # compile. Writes serialise naturally because each # `write_skill_file` depends on accumulated reads; the model has # no reason to issue parallel writes to the same path. - model_settings=ModelSettings(parallel_tool_calls=True), + model_settings=ModelSettings( + parallel_tool_calls=True, + extra_headers=get_extra_headers() or None, + ), ) diff --git a/openkb/skill/evaluator.py b/openkb/skill/evaluator.py index 426d4be80..9150ea541 100644 --- a/openkb/skill/evaluator.py +++ b/openkb/skill/evaluator.py @@ -44,7 +44,9 @@ from agents import Agent, Runner from agents.exceptions import MaxTurnsExceeded +from agents.model_settings import ModelSettings +from openkb.config import get_extra_headers from openkb.skill import extract_body, extract_frontmatter @@ -241,6 +243,7 @@ async def generate_eval_set( name="eval-set-generator", instructions=instructions, model=f"litellm/{model}", + model_settings=ModelSettings(extra_headers=get_extra_headers() or None), ) try: result = await Runner.run(agent, "Generate the eval set now.", max_turns=3) @@ -299,6 +302,7 @@ async def grade_one( name="trigger-grader", instructions=instructions, model=f"litellm/{model}", + model_settings=ModelSettings(extra_headers=get_extra_headers() or None), ) try: result = await Runner.run(agent, f"Question: {question}", max_turns=2) @@ -347,6 +351,7 @@ async def grade_coverage( name="coverage-grader", instructions=instructions, model=f"litellm/{model}", + model_settings=ModelSettings(extra_headers=get_extra_headers() or None), ) try: result = await Runner.run(agent, f"Question: {question}", max_turns=2) diff --git a/tests/conftest.py b/tests/conftest.py index 8d2271d34..3c5565510 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,14 @@ import pytest +@pytest.fixture(autouse=True) +def _reset_extra_headers(): + """Keep the process-wide LLM extra-headers stash from leaking across tests.""" + from openkb.config import set_extra_headers + yield + set_extra_headers({}) + + @pytest.fixture def kb_dir(tmp_path): """Create a minimal knowledge base directory structure for testing.""" diff --git a/tests/test_cli.py b/tests/test_cli.py index 65463566c..a94963ffc 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,6 +1,7 @@ import json from unittest.mock import patch +import pytest import yaml from click.testing import CliRunner @@ -366,3 +367,81 @@ async def fake_run_query(*_args, **_kwargs): assert "rnn" in saved assert "[[concepts/multi-head-attention]]" not in saved assert "multi head attention" in saved + + +class TestSetupLlmKey: + """_setup_llm_key: OAuth-provider warning skip + extra-headers stash.""" + + @staticmethod + def _make_kb(tmp_path, model, extra_headers=None): + openkb_dir = tmp_path / ".openkb" + openkb_dir.mkdir() + config = {"model": model} + if extra_headers is not None: + config["extra_headers"] = extra_headers + (openkb_dir / "config.yaml").write_text( + yaml.safe_dump(config), encoding="utf-8" + ) + return tmp_path + + @pytest.fixture(autouse=True) + def _clean_env(self, tmp_path, monkeypatch): + # Don't pick up the developer's real keys or global .env. + import openkb.config as config_mod + from openkb.cli import _KNOWN_PROVIDER_KEYS + + monkeypatch.setattr( + config_mod, "GLOBAL_CONFIG_DIR", tmp_path / "no-global" + ) + for key in ( + "LLM_API_KEY", + "GITHUB_COPILOT_API_KEY", + "CHATGPT_API_KEY", + *_KNOWN_PROVIDER_KEYS, + ): + monkeypatch.delenv(key, raising=False) + + @pytest.mark.parametrize("model", [ + "github_copilot/gpt-5-mini", + "chatgpt/gpt-5.4", + ]) + def test_no_warning_for_oauth_providers(self, tmp_path, capsys, model): + from openkb.cli import _setup_llm_key + + kb = self._make_kb(tmp_path, model) + _setup_llm_key(kb) + assert "No LLM API key found" not in capsys.readouterr().out + + def test_warning_for_api_key_provider_without_key(self, tmp_path, capsys): + from openkb.cli import _setup_llm_key + + kb = self._make_kb(tmp_path, "gpt-5.4-mini") + _setup_llm_key(kb) + assert "No LLM API key found" in capsys.readouterr().out + + def test_extra_headers_stashed_from_config(self, tmp_path): + from openkb.cli import _setup_llm_key + from openkb.config import get_extra_headers + + kb = self._make_kb( + tmp_path, + "github_copilot/gpt-5-mini", + extra_headers={ + "Editor-Version": "vscode/1.95.0", + "Copilot-Integration-Id": "vscode-chat", + }, + ) + _setup_llm_key(kb) + assert get_extra_headers() == { + "Editor-Version": "vscode/1.95.0", + "Copilot-Integration-Id": "vscode-chat", + } + + def test_extra_headers_reset_when_config_has_none(self, tmp_path): + from openkb.cli import _setup_llm_key + from openkb.config import get_extra_headers, set_extra_headers + + set_extra_headers({"Stale": "1"}) + kb = self._make_kb(tmp_path, "gpt-5.4-mini") + _setup_llm_key(kb) + assert get_extra_headers() == {} diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 84b4ff768..adbae6766 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -2046,3 +2046,54 @@ def test_plan_prompt_keeps_topic_itself_guard(): assert "just the document topic itself" in _CONCEPTS_PLAN_USER + + +class TestLLMCallExtraHeaders: + """Config-driven extra headers reach the litellm calls (issue #93).""" + + def test_llm_call_injects_extra_headers(self): + from openkb.agent.compiler import _llm_call + from openkb.config import set_extra_headers + + set_extra_headers({"Editor-Version": "vscode/1.95.0"}) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion(["ok"])) + out = _llm_call("m", [{"role": "user", "content": "hi"}], "step") + assert out == "ok" + kwargs = mock_litellm.completion.call_args.kwargs + assert kwargs["extra_headers"] == {"Editor-Version": "vscode/1.95.0"} + + def test_llm_call_no_extra_headers_by_default(self): + from openkb.agent.compiler import _llm_call + + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion(["ok"])) + _llm_call("m", [{"role": "user", "content": "hi"}], "step") + assert "extra_headers" not in mock_litellm.completion.call_args.kwargs + + def test_llm_call_explicit_kwarg_wins_over_config(self): + from openkb.agent.compiler import _llm_call + from openkb.config import set_extra_headers + + set_extra_headers({"Editor-Version": "from-config"}) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion(["ok"])) + _llm_call( + "m", [{"role": "user", "content": "hi"}], "step", + extra_headers={"Editor-Version": "explicit"}, + ) + kwargs = mock_litellm.completion.call_args.kwargs + assert kwargs["extra_headers"] == {"Editor-Version": "explicit"} + + @pytest.mark.asyncio + async def test_llm_call_async_injects_extra_headers(self): + from openkb.agent.compiler import _llm_call_async + from openkb.config import set_extra_headers + + set_extra_headers({"Copilot-Integration-Id": "vscode-chat"}) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.acompletion = AsyncMock(side_effect=_mock_acompletion(["ok"])) + out = await _llm_call_async("m", [{"role": "user", "content": "hi"}], "step") + assert out == "ok" + kwargs = mock_litellm.acompletion.call_args.kwargs + assert kwargs["extra_headers"] == {"Copilot-Integration-Id": "vscode-chat"} diff --git a/tests/test_config.py b/tests/test_config.py index 35704a6bf..22b720f77 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,4 +1,11 @@ -from openkb.config import DEFAULT_CONFIG, load_config, save_config +from openkb.config import ( + DEFAULT_CONFIG, + get_extra_headers, + load_config, + resolve_extra_headers, + save_config, + set_extra_headers, +) def test_default_config_keys(): @@ -45,3 +52,53 @@ def test_load_overrides_defaults(tmp_path): assert loaded["pageindex_threshold"] == 100 # Non-overridden defaults still present assert loaded["language"] == "en" + + +# --- extra_headers ----------------------------------------------------------- + +def test_resolve_extra_headers_absent_returns_empty(): + assert resolve_extra_headers({}) == {} + + +def test_resolve_extra_headers_valid_mapping(): + config = {"extra_headers": { + "Editor-Version": "vscode/1.95.0", + "Copilot-Integration-Id": "vscode-chat", + }} + assert resolve_extra_headers(config) == { + "Editor-Version": "vscode/1.95.0", + "Copilot-Integration-Id": "vscode-chat", + } + + +def test_resolve_extra_headers_stringifies_scalar_values(): + # YAML may parse version-ish values as numbers. + config = {"extra_headers": {"X-Api-Version": 2024, "X-Ratio": 1.5}} + assert resolve_extra_headers(config) == {"X-Api-Version": "2024", "X-Ratio": "1.5"} + + +def test_resolve_extra_headers_non_mapping_ignored(): + assert resolve_extra_headers({"extra_headers": ["Editor-Version: x"]}) == {} + assert resolve_extra_headers({"extra_headers": "Editor-Version: x"}) == {} + + +def test_resolve_extra_headers_skips_bad_entries(): + config = {"extra_headers": { + "Good": "value", + "": "empty-key-skipped", + "NoneValue": None, + "ListValue": ["a"], + 123: "non-string-key-skipped", + }} + assert resolve_extra_headers(config) == {"Good": "value"} + + +def test_extra_headers_stash_roundtrip_and_isolation(): + set_extra_headers({"A": "1"}) + got = get_extra_headers() + assert got == {"A": "1"} + # Mutating the returned copy must not affect the stash. + got["B"] = "2" + assert get_extra_headers() == {"A": "1"} + set_extra_headers({}) + assert get_extra_headers() == {} diff --git a/tests/test_query.py b/tests/test_query.py index e9585d327..4b19b17dc 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -144,3 +144,22 @@ def fake_print(*_args, **_kwargs): assert called["count"] == 1 assert fake_tty.getvalue() == "" + + +class TestQueryAgentExtraHeaders: + """Config-driven extra headers reach the agents-SDK model settings.""" + + def test_extra_headers_applied_from_stash(self, tmp_path): + from openkb.config import set_extra_headers + + set_extra_headers({"Editor-Version": "vscode/1.95.0"}) + agent = build_query_agent(str(tmp_path), "github_copilot/gpt-5-mini") + assert agent.model_settings.extra_headers == { + "Editor-Version": "vscode/1.95.0" + } + # Existing settings are preserved. + assert agent.model_settings.parallel_tool_calls is False + + def test_no_extra_headers_by_default(self, tmp_path): + agent = build_query_agent(str(tmp_path), "gpt-4o-mini") + assert agent.model_settings.extra_headers is None