From 820b989359cca0fb245c34e99f2168134e7aa8db Mon Sep 17 00:00:00 2001 From: Wachiravit Thitagarn Date: Wed, 29 Jul 2026 16:35:37 +0700 Subject: [PATCH] fix: use resolved LLM key for PageIndex --- openkb/cli.py | 7 +++- openkb/indexer.py | 72 ++++++++++++++++++++++++++++++++++----- tests/test_add_command.py | 12 +++++-- tests/test_indexer.py | 71 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 149 insertions(+), 13 deletions(-) diff --git a/openkb/cli.py b/openkb/cli.py index c9e54318..97dc15d5 100644 --- a/openkb/cli.py +++ b/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) diff --git a/openkb/indexer.py b/openkb/indexer.py index 6a4ae1ee..00cfa1aa 100644 --- a/openkb/indexer.py +++ b/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() diff --git a/tests/test_add_command.py b/tests/test_add_command.py index 3f51788e..1d5be4d1 100644 --- a/tests/test_add_command.py +++ b/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": []}) diff --git a/tests/test_indexer.py b/tests/test_indexer.py index e0843fa8..af6810a1 100644 --- a/tests/test_indexer.py +++ b/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)