fix(cli): only propagate LLM_API_KEY to the active provider key (#77)

Author: KylinMountain

openkb/cli.py

@@ -125,17 +125,19 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None:
     else:
         litellm.api_key = api_key
 
-        # Dynamically set the provider-specific env var when possible
         if provider:
+            # Active provider is known — set only its key, so LLM_API_KEY is not
+            # sprayed into unrelated provider keys (e.g. MISTRAL_API_KEY, which the
+            # Mistral parser treats as a real Mistral credential).
             provider_env = f"{provider.upper()}_API_KEY"
             if not os.environ.get(provider_env):
                 os.environ[provider_env] = api_key
-
-        # Fallback: also set common provider keys so multi-provider
-        # configs (e.g. PageIndex Cloud) still work
-        for env_var in _KNOWN_PROVIDER_KEYS:
-            if not os.environ.get(env_var):
-                os.environ[env_var] = api_key
+        else:
+            # Provider couldn't be determined — fall back to setting the common
+            # provider keys so multi-provider configs still work.
+            for env_var in _KNOWN_PROVIDER_KEYS:
+                if not os.environ.get(env_var):
+                    os.environ[env_var] = api_key
 
 # Supported document extensions for the `add` command
 SUPPORTED_EXTENSIONS = {

tests/test_cli.py

@@ -365,3 +365,17 @@ 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
+
+
+def test_setup_llm_key_does_not_spray_unrelated_provider_keys(tmp_path, monkeypatch):
+    import os
+    from openkb.cli import _setup_llm_key
+    # KB with an openai model (known provider)
+    openkb_dir = tmp_path / ".openkb"; openkb_dir.mkdir()
+    (openkb_dir / "config.yaml").write_text("model: openai/gpt-4o\n")
+    for k in ("MISTRAL_API_KEY", "OPENAI_API_KEY", "ANTHROPIC_API_KEY"):
+        monkeypatch.delenv(k, raising=False)
+    monkeypatch.setenv("LLM_API_KEY", "sk-test")
+    _setup_llm_key(tmp_path)
+    assert os.environ.get("OPENAI_API_KEY") == "sk-test"   # active provider set
+    assert os.environ.get("MISTRAL_API_KEY") is None        # unrelated provider NOT sprayed