From 62877eef28428098f346dcbf98b757be96fa0724 Mon Sep 17 00:00:00 2001 From: Wachiravit Thitagarn Date: Wed, 29 Jul 2026 16:35:37 +0700 Subject: [PATCH] fix: use latest UI LLM settings for ingest --- openkb/api.py | 5 ++- openkb/api_helpers.py | 14 +++----- openkb/cli.py | 16 +++++++-- openkb/indexer.py | 53 +++++++++++++++++++++++++--- openkb/watch_service.py | 9 +---- tests/test_add_command.py | 32 +++++++++++++++++ tests/test_indexer.py | 70 +++++++++++++++++++++++++++++++++++++ tests/test_watch_service.py | 17 +++++++++ 8 files changed, 188 insertions(+), 28 deletions(-) diff --git a/openkb/api.py b/openkb/api.py index 16b5db87..739d10cd 100644 --- a/openkb/api.py +++ b/openkb/api.py @@ -248,7 +248,6 @@ async def add_endpoint( _: None = Depends(require_bearer_token), ) -> Any: resolved_kb_dir = _resolve_kb(kb) - bundle = resolve_credential_bundle(resolved_kb_dir) if not files: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -263,10 +262,10 @@ async def add_endpoint( saved_uploads = await _write_add_uploads(reserved, files) if _parse_stream_form(stream): return StreamingResponse( - _stream_add_uploads(kb, resolved_kb_dir, saved_uploads, bundle=bundle), + _stream_add_uploads(kb, resolved_kb_dir, saved_uploads), media_type="text/event-stream", ) - return await _run_add_uploads(kb, resolved_kb_dir, saved_uploads, bundle=bundle) + return await _run_add_uploads(kb, resolved_kb_dir, saved_uploads) @app.post("/api/v1/query", response_model=QueryResponse) async def query_endpoint( diff --git a/openkb/api_helpers.py b/openkb/api_helpers.py index d42bf4a9..bda42d0d 100644 --- a/openkb/api_helpers.py +++ b/openkb/api_helpers.py @@ -335,12 +335,10 @@ async def _run_add_uploads( kb: str, kb_dir: Path, saved_uploads: list[tuple[Path, str]], - *, - bundle=None, ) -> AddResponse: results = [] for saved_path, original_name in saved_uploads: - results.append(await _add_saved_file(kb_dir, saved_path, original_name, bundle=bundle)) + results.append(await _add_saved_file(kb_dir, saved_path, original_name)) return _summarize_add_results(kb, results) @@ -348,8 +346,6 @@ async def _stream_add_uploads( kb: str, kb_dir: Path, saved_uploads: list[tuple[Path, str]], - *, - bundle=None, ) -> AsyncIterator[str]: yield _sse( "start", @@ -366,7 +362,7 @@ async def _stream_add_uploads( "file_start", {"original_name": original_name, "saved_path": str(saved_path)}, ) - item = await _add_saved_file(kb_dir, saved_path, original_name, bundle=bundle) + item = await _add_saved_file(kb_dir, saved_path, original_name) results.append(item) yield _sse("file_done", _model_payload(item)) final = _summarize_add_results(kb, results) @@ -378,10 +374,8 @@ async def _stream_add_uploads( yield _sse("done", {}) -async def _add_saved_file( - kb_dir: Path, saved_path: Path, original_name: str, *, bundle=None -) -> AddFileItem: - result = await run_in_threadpool(_add_for_api, saved_path, kb_dir, bundle=bundle) +async def _add_saved_file(kb_dir: Path, saved_path: Path, original_name: str) -> AddFileItem: + result = await run_in_threadpool(_add_for_api, saved_path, kb_dir) item = AddFileItem(**result.__dict__) item.original_name = original_name if item.status == "skipped": diff --git a/openkb/cli.py b/openkb/cli.py index c9e54318..8560c2b1 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) @@ -647,7 +652,7 @@ class AddFileResult: message: str -def _add_for_api(file_path: Path, kb_dir: Path, *, bundle=None) -> AddFileResult: +def _add_for_api(file_path: Path, kb_dir: Path) -> AddFileResult: """Run the locked add pipeline and return a structured result for the API. Reuses the upstream ``add_single_file`` (which already holds the ingest @@ -656,6 +661,13 @@ def _add_for_api(file_path: Path, kb_dir: Path, *, bundle=None) -> AddFileResult ``AddFileResult``; on ``skipped`` the caller (api._add_saved_file) deletes the freshly uploaded raw copy to avoid orphaning it. """ + # Resolve at the start of EACH file ingest, not when an upload request or + # long-lived watcher starts. Settings can be changed from the Workbench + # while either is active; a captured bundle would keep using the old model, + # key, base URL, headers, and timeout until the next request/restart. + from openkb.config import resolve_credential_bundle + + bundle = resolve_credential_bundle(kb_dir) status_str = add_single_file(file_path, kb_dir, bundle=bundle) if status_str == "skipped": message = f"Already in knowledge base: {file_path.name}" diff --git a/openkb/indexer.py b/openkb/indexer.py index 6a4ae1ee..2d3e6770 100644 --- a/openkb/indexer.py +++ b/openkb/indexer.py @@ -11,12 +11,27 @@ 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 client skips + that premature check; LiteLLM still validates the per-call key/base URL when + the indexing request is made. + """ + + @staticmethod + def _validate_llm_provider(model: str) -> None: + return None + + @dataclass class IndexResult: """Result of indexing a long document via PageIndex.""" @@ -153,7 +168,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 +196,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,9 +231,14 @@ 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( + client_type = ( + _PerRequestPageIndexClient + if bundle is not None and bundle.api_key and not pageindex_api_key + else PageIndexClient + ) + client = client_type( api_key=pageindex_api_key or None, model=model, storage_path=str(openkb_dir), diff --git a/openkb/watch_service.py b/openkb/watch_service.py index 1748a200..c76cf3c0 100644 --- a/openkb/watch_service.py +++ b/openkb/watch_service.py @@ -31,7 +31,6 @@ from watchdog.observers import Observer from openkb.cli import SUPPORTED_EXTENSIONS, _add_for_api -from openkb.config import LlmCredentialBundle, resolve_credential_bundle from openkb.watcher import start_watch # How many recent events to retain per KB for status() and SSE replay. @@ -55,7 +54,6 @@ class WatcherState: counters: dict[str, int] = field( default_factory=lambda: {"added": 0, "skipped": 0, "failed": 0} ) - bundle: LlmCredentialBundle | None = None _seq: int = 0 _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) @@ -130,7 +128,7 @@ def _process_file(state: WatcherState, raw_path: str) -> None: }, ) try: - result = _add_for_api(path, state.kb_dir, bundle=state.bundle) + result = _add_for_api(path, state.kb_dir) except Exception as exc: # worker must never die _record_event( state, @@ -184,16 +182,11 @@ def start(self, kb: str, kb_dir: Path, debounce: float = 2.0) -> WatcherState: self._watchers.pop(kb, None) raw_dir = kb_dir / "raw" raw_dir.mkdir(parents=True, exist_ok=True) - # Resolve the KB's credentials once at watcher start so background - # ingests use the same per-KB bundle as REST endpoints instead of - # mutating process-wide env state. - bundle = resolve_credential_bundle(kb_dir) state = WatcherState( kb=kb, kb_dir=kb_dir, raw_dir=raw_dir, debounce=debounce, - bundle=bundle, started_at=time.time(), events=deque(maxlen=self._max_events), ) diff --git a/tests/test_add_command.py b/tests/test_add_command.py index 3f51788e..369095fd 100644 --- a/tests/test_add_command.py +++ b/tests/test_add_command.py @@ -41,6 +41,38 @@ def test_returns_none_if_no_openkb(self, tmp_path, monkeypatch): assert result is None +def test_add_for_api_resolves_latest_persisted_credentials_each_call(tmp_path): + from openkb.cli import _add_for_api + + (tmp_path / ".openkb").mkdir() + (tmp_path / ".openkb" / "config.yaml").write_text("model: openai/cs/abc_ex\n", encoding="utf-8") + env_path = tmp_path / ".env" + env_path.write_text( + "LLM_API_KEY=old-key\nOPENAI_API_BASE=https://old.example/v1\n", + encoding="utf-8", + ) + doc = tmp_path / "doc.md" + doc.write_text("# Doc\n", encoding="utf-8") + seen = [] + + def fake_add(_path, _kb_dir, *, bundle=None): + seen.append((bundle.api_key, bundle.base_url)) + return "added" + + with patch("openkb.cli.add_single_file", side_effect=fake_add): + _add_for_api(doc, tmp_path) + env_path.write_text( + "LLM_API_KEY=new-key\nOPENAI_API_BASE=https://new.example/v1\n", + encoding="utf-8", + ) + _add_for_api(doc, tmp_path) + + assert seen == [ + ("old-key", "https://old.example/v1"), + ("new-key", "https://new.example/v1"), + ] + + class TestAddCommand: def _setup_kb(self, tmp_path): """Create a minimal KB structure.""" diff --git a/tests/test_indexer.py b/tests/test_indexer.py index e0843fa8..9c27a4b4 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,57 @@ 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["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/cs/abc_ex\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) diff --git a/tests/test_watch_service.py b/tests/test_watch_service.py index 27be483f..680dff39 100644 --- a/tests/test_watch_service.py +++ b/tests/test_watch_service.py @@ -14,6 +14,7 @@ from pathlib import Path from openkb.cli import AddFileResult +from openkb.config import LlmCredentialBundle from openkb.watch_service import ( WatcherState, WatchRegistry, @@ -117,6 +118,22 @@ def test_worker_added_records_file_start_done_and_counter(kb_dir, monkeypatch): assert state.counters == {"added": 1, "skipped": 0, "failed": 0} +def test_worker_does_not_reuse_bundle_captured_at_watcher_start(kb_dir, monkeypatch): + state = _make_state(kb_dir) + state.bundle = LlmCredentialBundle(api_key="stale-key") + state.queue.put([str(kb_dir / "raw" / "paper.md")]) + seen = [] + + def fake_add(path, kb, bundle=None): + seen.append(bundle) + return AddFileResult(path.name, str(path), "added", "Added.") + + monkeypatch.setattr("openkb.watch_service._add_for_api", fake_add) + _drain_worker(state) + + assert seen == [None] + + def test_worker_skipped_and_failed_branches(kb_dir, monkeypatch): state = _make_state(kb_dir) state.queue.put([str(kb_dir / "raw" / "dup.md"), str(kb_dir / "raw" / "boom.md")])