Skip to content

fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wachiravit-thitagran wants to merge 1 commit into
base:main
Choose a base branch
from
Draft

fix: use resolved LLM key for PageIndex #209

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 6 additions & 1 deletion openkb/cli.py
View file
Open in desktop
Original file line numberDiff line numberDiff 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
72 changes: 63 additions & 9 deletions openkb/indexer.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down
12 changes: 9 additions & 3 deletions tests/test_add_command.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down
71 changes: 71 additions & 0 deletions tests/test_indexer.py
View file
Open in desktop
Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
fix: use resolved LLM key for PageIndex by wachiravit-thitagran · Pull Request #209 · VectifyAI/OpenKB · GitHub
Skip to content
fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

wachiravit-thitagran
wants to merge
1
commit into

base:main
Choose a base branch

from

 Draft


fix: use resolved LLM key for PageIndex
#209
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


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


7 changes: 6 additions & 1 deletion
openkb/cli.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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

72 changes: 63 additions & 9 deletions
openkb/indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down

12 changes: 9 additions & 3 deletions
tests/test_add_command.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down

71 changes: 71 additions & 0 deletions
tests/test_indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix: use resolved LLM key for PageIndex by wachiravit-thitagran · Pull Request #209 · VectifyAI/OpenKB · GitHub
Skip to content
fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

wachiravit-thitagran
wants to merge
1
commit into

base:main
Choose a base branch

from

 Draft


fix: use resolved LLM key for PageIndex
#209
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


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


7 changes: 6 additions & 1 deletion
openkb/cli.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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

72 changes: 63 additions & 9 deletions
openkb/indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down

12 changes: 9 additions & 3 deletions
tests/test_add_command.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down

71 changes: 71 additions & 0 deletions
tests/test_indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix: use resolved LLM key for PageIndex by wachiravit-thitagran · Pull Request #209 · VectifyAI/OpenKB · GitHub
Skip to content
fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

wachiravit-thitagran
wants to merge
1
commit into

base:main
Choose a base branch

from

 Draft


fix: use resolved LLM key for PageIndex
#209
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


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


7 changes: 6 additions & 1 deletion
openkb/cli.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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

72 changes: 63 additions & 9 deletions
openkb/indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down

12 changes: 9 additions & 3 deletions
tests/test_add_command.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down

71 changes: 71 additions & 0 deletions
tests/test_indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
fix: use resolved LLM key for PageIndex by wachiravit-thitagran · Pull Request #209 · VectifyAI/OpenKB · GitHub
Skip to content
fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

wachiravit-thitagran
wants to merge
1
commit into

base:main
Choose a base branch

from

 Draft


fix: use resolved LLM key for PageIndex
#209
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


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


7 changes: 6 additions & 1 deletion
openkb/cli.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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

72 changes: 63 additions & 9 deletions
openkb/indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down

12 changes: 9 additions & 3 deletions
tests/test_add_command.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down

71 changes: 71 additions & 0 deletions
tests/test_indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix: use resolved LLM key for PageIndex by wachiravit-thitagran · Pull Request #209 · VectifyAI/OpenKB · GitHub
Skip to content
fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

wachiravit-thitagran
wants to merge
1
commit into

base:main
Choose a base branch

from

 Draft


fix: use resolved LLM key for PageIndex
#209
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


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


7 changes: 6 additions & 1 deletion
openkb/cli.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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

72 changes: 63 additions & 9 deletions
openkb/indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down

12 changes: 9 additions & 3 deletions
tests/test_add_command.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down

71 changes: 71 additions & 0 deletions
tests/test_indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
fix: use resolved LLM key for PageIndex by wachiravit-thitagran · Pull Request #209 · VectifyAI/OpenKB · GitHub
Skip to content
fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

wachiravit-thitagran
wants to merge
1
commit into

base:main
Choose a base branch

from

 Draft


fix: use resolved LLM key for PageIndex
#209
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


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


7 changes: 6 additions & 1 deletion
openkb/cli.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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

72 changes: 63 additions & 9 deletions
openkb/indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down

12 changes: 9 additions & 3 deletions
tests/test_add_command.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down

71 changes: 71 additions & 0 deletions
tests/test_indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
fix: use resolved LLM key for PageIndex by wachiravit-thitagran · Pull Request #209 · VectifyAI/OpenKB · GitHub
Skip to content
fix: use resolved LLM key for PageIndex#209

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

wachiravit-thitagran
wants to merge
1
commit into

base:main
Choose a base branch

from

 Draft


fix: use resolved LLM key for PageIndex
#209
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


1 commit


Select commit


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


7 changes: 6 additions & 1 deletion
openkb/cli.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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

72 changes: 63 additions & 9 deletions
openkb/indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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."""
Expand DownExpand Up@@ -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
Expand All@@ -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
Expand All@@ -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()

Expand Down

12 changes: 9 additions & 3 deletions
tests/test_add_command.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand All@@ -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")
Expand All@@ -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
Expand All@@ -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": []})

Expand Down

71 changes: 71 additions & 0 deletions
tests/test_indexer.py

View file


Open in desktop

Original file line numberDiff line numberDiff 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 DownExpand 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 DownExpand Up@@ -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)
Expand Down