fix: use latest UI LLM settings for ingest

Author: wachiravit-thitagran

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(

openkb/api_helpers.py

@@ -335,21 +335,17 @@ 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)
 
 
 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":

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}"

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),

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),
             )

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."""

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)

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")])