fix: use resolved LLM key for PageIndex

openkb/cli.py

@@ -527,7 +527,12 @@ def commit_body(snapshot) -> None:
             try:
                 from openkb.indexer import index_long_document
 
-                index_result = index_long_document(result.raw_path, kb_dir, doc_name=doc_name)
+                index_result = index_long_document(
+                    result.raw_path,
+                    kb_dir,
+                    doc_name=doc_name,
+                    bundle=bundle,
+                )
             except Exception as exc:
                 click.echo(f"  [ERROR] Indexing failed: {exc}")
                 logger.debug("Indexing traceback:", exc_info=True)

openkb/indexer.py

@@ -11,12 +11,32 @@
 
 from pageindex import IndexConfig, PageIndexClient
 
-from openkb.config import resolve_concurrency, resolve_effective_config
+from openkb.config import LlmCredentialBundle, resolve_concurrency, resolve_effective_config
 from openkb.tree_renderer import render_summary_md
 
 logger = logging.getLogger(__name__)
 
 
+class _PerRequestPageIndexClient(PageIndexClient):
+    """PageIndex local client for credentials supplied in ``IndexConfig``.
+
+    PageIndex validates only process-wide provider environment variables during
+    construction, before its scoped ``llm_params`` reach LiteLLM. REST requests
+    deliberately keep credentials out of process globals, so this adapter
+    accepts the explicitly resolved ``LLM_API_KEY``. Without that key, normal
+    PageIndex provider validation still applies.
+    """
+
+    def __init__(self, *, llm_api_key: str | None = None, **kwargs: Any) -> None:
+        self._llm_api_key = llm_api_key
+        super().__init__(**kwargs)
+
+    def _validate_llm_provider(self, model: str) -> None:
+        if self._llm_api_key:
+            return
+        super()._validate_llm_provider(model)
+
+
 @dataclass
 class IndexResult:
     """Result of indexing a long document via PageIndex."""
@@ -153,7 +173,11 @@ def _write_long_doc_artifacts(
     return summary_path
 
 
-def _build_index_config(config: dict[str, Any]) -> IndexConfig:
+def _build_index_config(
+    config: dict[str, Any],
+    *,
+    bundle: LlmCredentialBundle | None = None,
+) -> IndexConfig:
     """Build the PageIndex ``IndexConfig`` for local indexing.
 
     Forwards the KB's ``concurrency`` setting to PageIndex, which caps how many
@@ -177,10 +201,29 @@ def _build_index_config(config: dict[str, Any]) -> IndexConfig:
                 "config: 'concurrency' is set but the installed PageIndex "
                 "version does not support it yet — ignoring it."
             )
+    if bundle is not None:
+        llm_params = {
+            key: value
+            for key, value in {
+                "api_key": bundle.api_key,
+                "base_url": bundle.base_url,
+                "extra_headers": bundle.extra_headers or None,
+                "timeout": bundle.timeout,
+            }.items()
+            if value is not None
+        }
+        if llm_params:
+            kwargs["llm_params"] = llm_params
     return IndexConfig(**kwargs)
 
 
-def index_long_document(pdf_path: Path, kb_dir: Path, doc_name: str | None = None) -> IndexResult:
+def index_long_document(
+    pdf_path: Path,
+    kb_dir: Path,
+    doc_name: str | None = None,
+    *,
+    bundle: LlmCredentialBundle | None = None,
+) -> IndexResult:
     """Index a long PDF document using PageIndex and write wiki pages.
 
     ``doc_name`` is the collision-resistant wiki name used for all written
@@ -193,13 +236,24 @@ def index_long_document(pdf_path: Path, kb_dir: Path, doc_name: str | None = Non
     model: str = config.get("model", "gpt-5.4")
     pageindex_api_key = os.environ.get("PAGEINDEX_API_KEY", "")
 
-    index_config = _build_index_config(config)
+    index_config = _build_index_config(config, bundle=bundle)
 
-    client = PageIndexClient(
-        api_key=pageindex_api_key or None,
-        model=model,
-        storage_path=str(openkb_dir),
-        index_config=index_config,
+    client_type = (
+        _PerRequestPageIndexClient
+        if bundle is not None and bundle.api_key and not pageindex_api_key
+        else PageIndexClient
+    )
+    client_kwargs: dict[str, Any] = {
+        "api_key": pageindex_api_key or None,
+        "model": model,
+        "storage_path": str(openkb_dir),
+        "index_config": index_config,
+    }
+    if client_type is _PerRequestPageIndexClient:
+        assert bundle is not None
+        client_kwargs["llm_api_key"] = bundle.api_key
+    client = client_type(
+        **client_kwargs,
     )
     col = client.collection()
 

tests/test_add_command.py

@@ -167,6 +167,7 @@ def test_long_doc_rollback_removes_only_the_new_blob(self, tmp_path):
         .openkb/files, while a pre-existing blob (another document) survives —
         the targeted track_new must not touch blobs this add didn't create."""
         from openkb.cli import add_single_file
+        from openkb.config import LlmCredentialBundle
         from openkb.indexer import IndexResult
 
         kb_dir = self._setup_kb(tmp_path)
@@ -176,8 +177,12 @@ def test_long_doc_rollback_removes_only_the_new_blob(self, tmp_path):
         other.write_bytes(b"another-doc-keep-me")
 
         new_id = "11111111-1111-1111-1111-111111111111"
+        bundle = LlmCredentialBundle(api_key="ui-key")
+        seen_bundle = None
 
-        def fake_index(raw_path, kb_dir_arg, doc_name=None):
+        def fake_index(raw_path, kb_dir_arg, doc_name=None, *, bundle=None):
+            nonlocal seen_bundle
+            seen_bundle = bundle
             (files / f"{new_id}.pdf").write_bytes(b"new-blob")
             (files / new_id / "images").mkdir(parents=True)
             (files / new_id / "images" / "p1.png").write_bytes(b"img")
@@ -194,9 +199,10 @@ def fake_index(raw_path, kb_dir_arg, doc_name=None):
             patch("openkb.cli.time.sleep"),
             patch("openkb.cli._setup_llm_key"),
         ):
-            outcome = add_single_file(doc, kb_dir)
+            outcome = add_single_file(doc, kb_dir, bundle=bundle)
 
         assert outcome == "failed"
+        assert seen_bundle is bundle
         assert not (files / f"{new_id}.pdf").exists()  # new blob rolled back
         assert not (files / new_id).exists()  # new images subtree rolled back
         assert other.read_bytes() == b"another-doc-keep-me"  # pre-existing survives
@@ -216,7 +222,7 @@ def test_long_doc_dedup_hit_does_not_delete_existing_blob(self, tmp_path):
         existing_blob = files / f"{existing_id}.pdf"
         existing_blob.write_bytes(b"pre-existing-do-not-delete")
 
-        def fake_index_dedup(raw_path, kb_dir_arg, doc_name=None):
+        def fake_index_dedup(raw_path, kb_dir_arg, doc_name=None, *, bundle=None):
             # Dedup hit: return the existing doc_id, create NO new blob.
             return IndexResult(doc_id=existing_id, description="", tree={"structure": []})
 

tests/test_indexer.py

@@ -6,7 +6,9 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from pageindex import PageIndexClient
 
+from openkb.config import LlmCredentialBundle
 from openkb.indexer import (
     IndexResult,
     _build_index_config,
@@ -94,6 +96,23 @@ def test_no_warning_when_supported(self, monkeypatch, caplog):
             _build_index_config({"concurrency": 8})
         assert caplog.text == ""
 
+    def test_forwards_rest_credentials_as_pageindex_per_call_params(self):
+        bundle = LlmCredentialBundle(
+            api_key="ui-key",
+            base_url="https://gateway.example/v1",
+            extra_headers={"X-Tenant": "alpha"},
+            timeout=45,
+        )
+
+        cfg = _build_index_config({}, bundle=bundle)
+
+        assert cfg.llm_params == {
+            "api_key": "ui-key",
+            "base_url": "https://gateway.example/v1",
+            "extra_headers": {"X-Tenant": "alpha"},
+            "timeout": 45,
+        }
+
 
 class TestNormalizePageContent:
     def test_normalizes_pageindex_dicts(self):
@@ -313,6 +332,58 @@ def test_concurrency_flows_from_kb_config(self, kb_dir, sample_tree, tmp_path):
         _, kwargs = mock_cls.call_args
         assert kwargs["index_config"].max_concurrency == 7
 
+    def test_rest_bundle_reaches_local_pageindex_client(self, kb_dir, sample_tree, tmp_path):
+        doc_id = "bundle-123"
+        fake_col = self._make_fake_collection(doc_id, sample_tree)
+        fake_client = MagicMock()
+        fake_client.collection.return_value = fake_col
+        bundle = LlmCredentialBundle(
+            api_key="ui-key",
+            base_url="https://gateway.example/v1",
+        )
+        pdf_path = tmp_path / "report.pdf"
+        pdf_path.write_bytes(b"%PDF-1.4 fake")
+
+        with (
+            patch(
+                "openkb.indexer._PerRequestPageIndexClient",
+                return_value=fake_client,
+            ) as mock_cls,
+            patch("openkb.images.convert_pdf_to_pages", return_value=self._fake_pages()),
+        ):
+            index_long_document(pdf_path, kb_dir, bundle=bundle)
+
+        _, kwargs = mock_cls.call_args
+        assert kwargs["llm_api_key"] == "ui-key"
+        assert kwargs["index_config"].llm_params == {
+            "api_key": "ui-key",
+            "base_url": "https://gateway.example/v1",
+        }
+
+    def test_rest_bundle_does_not_require_provider_key_in_process_env(
+        self, kb_dir, sample_tree, tmp_path, monkeypatch
+    ):
+        monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+        monkeypatch.setattr("litellm.api_key", None)
+        (kb_dir / ".openkb" / "config.yaml").write_text(
+            "model: openai/gpt-5.4-pool\n", encoding="utf-8"
+        )
+        fake_col = self._make_fake_collection("bundle-456", sample_tree)
+        bundle = LlmCredentialBundle(
+            api_key="ui-key",
+            base_url="https://gateway.example/v1",
+        )
+        pdf_path = tmp_path / "report.pdf"
+        pdf_path.write_bytes(b"%PDF-1.4 fake")
+
+        with (
+            patch.object(PageIndexClient, "collection", return_value=fake_col),
+            patch("openkb.images.convert_pdf_to_pages", return_value=self._fake_pages()),
+        ):
+            result = index_long_document(pdf_path, kb_dir, bundle=bundle)
+
+        assert result.doc_id == "bundle-456"
+
     def test_cloud_page_content_is_normalized(self, kb_dir, sample_tree, tmp_path, monkeypatch):
         doc_id = "cloud-123"
         fake_col = self._make_fake_collection(doc_id, sample_tree)