Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions openkb/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(
Expand Down
14 changes: 4 additions & 10 deletions openkb/api_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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)
Expand All @@ -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":
Expand Down
16 changes: 14 additions & 2 deletions openkb/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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}"
Expand Down
53 changes: 48 additions & 5 deletions openkb/indexer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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),
Expand Down
9 changes: 1 addition & 8 deletions openkb/watch_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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),
)
Expand Down
32 changes: 32 additions & 0 deletions tests/test_add_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
70 changes: 70 additions & 0 deletions tests/test_indexer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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)
Expand Down
17 changes: 17 additions & 0 deletions tests/test_watch_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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")])
Expand Down