diff --git a/config.yaml.example b/config.yaml.example index eebf3bf6..93fad530 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -2,6 +2,12 @@ model: gpt-5.4 # LLM model (any LiteLLM-supported provider) language: en # Wiki output language pageindex_threshold: 20 # PDF pages threshold for PageIndex +# Optional: cap concurrent LLM calls during ingest (PageIndex indexing and +# concept/entity compilation — they never overlap, so one setting covers +# both). Lower it if you hit provider rate limits or "too many open files" on +# large PDFs. Omit to let each stage apply its own default. +# concurrency: 5 + # Optional: override the entity-type vocabulary used for entity pages. # Omit this key to use the default 7 types # (person, organization, place, product, work, event, other). diff --git a/examples/configuration/README.md b/examples/configuration/README.md index 58cb7d4d..6790cc80 100644 --- a/examples/configuration/README.md +++ b/examples/configuration/README.md @@ -70,6 +70,12 @@ model: gpt-5.4 # LLM model (any LiteLLM-supported provider) language: en # Wiki output language pageindex_threshold: 20 # PDF pages threshold for PageIndex +# Optional: cap concurrent LLM calls during ingest (PageIndex indexing and +# concept/entity compilation — they never overlap, so one setting covers +# both). Lower it if you hit provider rate limits or "too many open files" on +# large PDFs. Omit to let each stage apply its own default. +# concurrency: 5 + # Optional: override the entity-type vocabulary used for entity pages. # Omit this key to use the default 7 types # (person, organization, place, product, work, event, other). @@ -95,6 +101,7 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex | `model` | `gpt-5.4` | LLM used for all compile/query/chat work. | | `language` | `en` | Language the wiki is written in. | | `pageindex_threshold` | `20` | PDFs with this many pages **or more** take the long-doc (PageIndex) path; shorter ones go through the short-doc path. See [`pageindex-cloud/`](../pageindex-cloud/). | +| `concurrency` | `null` | Caps concurrent LLM calls OpenKB makes during ingest — both PageIndex's indexing of a long document and OpenKB's own concept/entity compilation. The two never run at once for the same document, so one setting covers both. Lower it if you hit provider rate limits or "too many open files" on large PDFs. `null` lets each stage apply its own default. | | `entity_types` | 7 defaults | Custom vocabulary for entity pages. `other` is always kept. | | `litellm:` | – | A pass-through block for LiteLLM. See below. | diff --git a/openkb/cli.py b/openkb/cli.py index f39d0b44..56ca1031 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -47,13 +47,14 @@ def filter(self, record: logging.LogRecord) -> bool: litellm.suppress_debug_info = True from dotenv import load_dotenv -from openkb.agent.compiler import compile_long_doc +from openkb.agent.compiler import DEFAULT_COMPILE_CONCURRENCY, compile_long_doc from openkb.config import ( DEFAULT_CONFIG, load_config, save_config, load_global_config, register_kb, + resolve_concurrency, resolve_extra_headers, set_extra_headers, resolve_timeout, @@ -555,6 +556,7 @@ def commit_body(snapshot) -> None: kb_dir, model, doc_description=index_result.description, + max_concurrency=resolve_concurrency(config) or DEFAULT_COMPILE_CONCURRENCY, ), label=f"Compiling long doc (doc_id={index_result.doc_id})", ) @@ -563,7 +565,13 @@ def commit_body(snapshot) -> None: raise RuntimeError(f"Converted document has no source artifact: {file_path.name}") source_path = result.source_path _run_compile_with_retry( - lambda: compile_short_doc(doc_name, source_path, kb_dir, model), + lambda: compile_short_doc( + doc_name, + source_path, + kb_dir, + model, + max_concurrency=resolve_concurrency(config) or DEFAULT_COMPILE_CONCURRENCY, + ), label="Compiling short doc", ) @@ -684,6 +692,7 @@ def commit_body(_snapshot) -> None: kb_dir, model, doc_description=cloud.description, + max_concurrency=resolve_concurrency(config) or DEFAULT_COMPILE_CONCURRENCY, ), label=f"Compiling imported doc (doc_id={doc_id})", ) @@ -1636,6 +1645,7 @@ def _classify(meta: dict) -> str: _setup_llm_key(kb_dir) config = load_config(openkb_dir / "config.yaml") model: str = config.get("model", DEFAULT_CONFIG["model"]) + max_concurrency = resolve_concurrency(config) or DEFAULT_COMPILE_CONCURRENCY # Import lazily and reference via the module so tests can patch # ``openkb.agent.compiler.compile_*`` and see the call. @@ -1671,7 +1681,16 @@ def _classify(meta: dict) -> str: click.echo(f"[{i}/{total}] Recompiling long doc {name}...") start = time.time() try: - asyncio.run(compiler.compile_long_doc(name, summary_path, doc_id, kb_dir, model)) + asyncio.run( + compiler.compile_long_doc( + name, + summary_path, + doc_id, + kb_dir, + model, + max_concurrency=max_concurrency, + ) + ) except Exception as exc: click.echo(f" [ERROR] Compilation failed: {exc}") logging.getLogger(__name__).debug("Recompile traceback:", exc_info=True) @@ -1691,7 +1710,15 @@ def _classify(meta: dict) -> str: click.echo(f"[{i}/{total}] Recompiling short doc {name}...") start = time.time() try: - asyncio.run(compiler.compile_short_doc(name, source_path, kb_dir, model)) + asyncio.run( + compiler.compile_short_doc( + name, + source_path, + kb_dir, + model, + max_concurrency=max_concurrency, + ) + ) except Exception as exc: click.echo(f" [ERROR] Compilation failed: {exc}") logging.getLogger(__name__).debug("Recompile traceback:", exc_info=True) diff --git a/openkb/config.py b/openkb/config.py index a5ee010b..99b69864 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -17,6 +17,13 @@ "model": "gpt-5.4", "language": "en", "pageindex_threshold": 20, + # Cap on concurrent LLM calls OpenKB makes during ingest — both PageIndex's + # indexing of a long document and OpenKB's own concept/entity compilation. + # These never run concurrently with each other for the same document, so + # one knob covers both. None = each stage applies its own built-in default. + # Lower it if you hit provider rate limits or "too many open files"; raise + # it to go faster. + "concurrency": None, } # Default entity-type vocabulary. Overridable per-KB via the optional @@ -176,6 +183,30 @@ def resolve_timeout(config: dict) -> float | None: return value +def resolve_concurrency(config: dict) -> int | None: + """Resolve the optional ``concurrency:`` key — the cap on concurrent LLM + calls OpenKB makes during ingest (PageIndex indexing and concept/entity + compilation alike; they never run at the same time for one document, so + one setting covers both). + + Returns ``None`` when absent, explicitly ``null``, or invalid (rejecting + bools and non-positive values with a warning when present but unusable — + an explicit ``null``/absent key is the normal "unset" case and stays + silent, matching ``resolve_timeout``). Callers apply their own default (or + omit the setting entirely) when this returns ``None``. + """ + value = config.get("concurrency") + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + logger.warning( + "config: 'concurrency' must be a positive integer, got %r — ignoring it.", + value, + ) + return None + return value + + def resolve_litellm_settings(config: dict) -> dict[str, Any]: """Resolve the optional ``litellm:`` mapping of LiteLLM module settings. diff --git a/openkb/indexer.py b/openkb/indexer.py index c8272200..3b55fc42 100644 --- a/openkb/indexer.py +++ b/openkb/indexer.py @@ -11,7 +11,7 @@ from pageindex import IndexConfig, PageIndexClient -from openkb.config import load_config +from openkb.config import load_config, resolve_concurrency from openkb.tree_renderer import render_summary_md logger = logging.getLogger(__name__) @@ -153,6 +153,33 @@ def _write_long_doc_artifacts( return summary_path +def _build_index_config(config: dict[str, Any]) -> IndexConfig: + """Build the PageIndex ``IndexConfig`` for local indexing. + + Forwards the KB's ``concurrency`` setting to PageIndex, which caps how many + indexing LLM calls run at once (guarding against "too many open files" fd + exhaustion on large documents). The value is only passed when set *and* the + installed PageIndex's ``IndexConfig`` declares the field, so OpenKB keeps + working against a pinned PageIndex that predates it (``IndexConfig`` + forbids unknown kwargs). + """ + kwargs: dict[str, Any] = { + "if_add_node_text": True, + "if_add_node_summary": True, + "if_add_doc_description": True, + } + concurrency = resolve_concurrency(config) + if concurrency is not None: + if "max_concurrency" in IndexConfig.model_fields: + kwargs["max_concurrency"] = concurrency + else: + logger.warning( + "config: 'concurrency' is set but the installed PageIndex " + "version does not support it yet — ignoring it." + ) + return IndexConfig(**kwargs) + + def index_long_document(pdf_path: Path, kb_dir: Path, doc_name: str | None = None) -> IndexResult: """Index a long PDF document using PageIndex and write wiki pages. @@ -166,11 +193,7 @@ 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 = IndexConfig( - if_add_node_text=True, - if_add_node_summary=True, - if_add_doc_description=True, - ) + index_config = _build_index_config(config) client = PageIndexClient( api_key=pageindex_api_key or None, diff --git a/pyproject.toml b/pyproject.toml index c9c79b4e..c59d948b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ keywords = ["ai", "rag", "retrieval", "knowledge-base", "llm", "pageindex", "age # returning empty Responses output (BerriAI/litellm#25429) and # auto-injects GitHub Copilot IDE-auth headers. dependencies = [ - "pageindex==0.3.0.dev1", + "pageindex==0.3.0.dev2", "markitdown[docx,pptx,xlsx,xls]==0.1.5", "trafilatura==2.0.0", "click==8.4.0", diff --git a/tests/test_add_command.py b/tests/test_add_command.py index ae84461e..3f51788e 100644 --- a/tests/test_add_command.py +++ b/tests/test_add_command.py @@ -97,6 +97,29 @@ def test_add_single_file_compile_failure_rolls_back_converted_artifacts(self, tm assert not (kb_dir / "wiki" / "sources" / "notes.md").exists() assert HashRegistry(kb_dir / ".openkb" / "hashes.json").all_entries() == {} + def test_add_forwards_concurrency_from_config(self, tmp_path): + from unittest.mock import AsyncMock + + from openkb.cli import add_single_file + + kb_dir = self._setup_kb(tmp_path) + (kb_dir / ".openkb" / "config.yaml").write_text( + "model: gpt-4o-mini\nconcurrency: 3\n", encoding="utf-8" + ) + doc = tmp_path / "notes.md" + doc.write_text("# Notes\n\nBody", encoding="utf-8") + + with ( + patch( + "openkb.agent.compiler.compile_short_doc", new_callable=AsyncMock + ) as mock_compile, + patch("openkb.cli._setup_llm_key"), + ): + outcome = add_single_file(doc, kb_dir) + + assert outcome == "added" + assert mock_compile.call_args.kwargs["max_concurrency"] == 3 + def test_add_single_file_uses_add_mutation_coordinator(self, tmp_path): from openkb.cli import add_single_file from openkb.converter import ConvertResult diff --git a/tests/test_config.py b/tests/test_config.py index be473fcd..90bbbcbb 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -5,6 +5,7 @@ get_extra_headers, get_timeout, load_config, + resolve_concurrency, resolve_extra_headers, resolve_litellm_settings, resolve_timeout, @@ -26,6 +27,53 @@ def test_default_config_values(): assert DEFAULT_CONFIG["pageindex_threshold"] == 20 +def test_concurrency_defaults_to_none(): + assert DEFAULT_CONFIG["concurrency"] is None + + +def test_load_concurrency_override(tmp_path): + config_path = tmp_path / "config.yaml" + config_path.write_text("concurrency: 12\n", encoding="utf-8") + assert load_config(config_path)["concurrency"] == 12 + + +def test_resolve_concurrency_absent_is_none(): + assert resolve_concurrency({}) is None + + +def test_resolve_concurrency_valid_value(): + assert resolve_concurrency({"concurrency": 3}) == 3 + + +def test_resolve_concurrency_rejects_bool(caplog): + with caplog.at_level(logging.WARNING, logger="openkb.config"): + result = resolve_concurrency({"concurrency": True}) + assert result is None + assert "concurrency" in caplog.text + + +def test_resolve_concurrency_rejects_non_positive(caplog): + with caplog.at_level(logging.WARNING, logger="openkb.config"): + assert resolve_concurrency({"concurrency": 0}) is None + assert "concurrency" in caplog.text + caplog.clear() + with caplog.at_level(logging.WARNING, logger="openkb.config"): + assert resolve_concurrency({"concurrency": -1}) is None + assert "concurrency" in caplog.text + + +def test_resolve_concurrency_rejects_non_int(): + assert resolve_concurrency({"concurrency": "3"}) is None + + +def test_resolve_concurrency_none_is_silent(caplog): + # Explicit null / absent is the normal "unset — caller applies its own + # default, or omits the setting entirely" case — no warning. + with caplog.at_level(logging.WARNING, logger="openkb.config"): + assert resolve_concurrency({"concurrency": None}) is None + assert caplog.text == "" + + def test_load_missing_file_returns_defaults(tmp_path): missing = tmp_path / "nonexistent" / "config.yaml" config = load_config(missing) diff --git a/tests/test_indexer.py b/tests/test_indexer.py index 1c9e4485..e0843fa8 100644 --- a/tests/test_indexer.py +++ b/tests/test_indexer.py @@ -2,11 +2,97 @@ from __future__ import annotations +import logging from unittest.mock import MagicMock, patch import pytest -from openkb.indexer import IndexResult, _normalize_page_content, index_long_document +from openkb.indexer import ( + IndexResult, + _build_index_config, + _normalize_page_content, + index_long_document, +) + + +class _FakeIndexConfigWithConcurrency: + """Stand-in for a PageIndex ``IndexConfig`` that declares ``max_concurrency``. + + Used instead of relying on whatever ``pageindex`` happens to be installed in + this environment, so the forwarding tests are deterministic regardless of + the currently-pinned PageIndex version (see ``test_forwards_...`` below). + """ + + model_fields = { + "if_add_node_text": None, + "if_add_node_summary": None, + "if_add_doc_description": None, + "max_concurrency": None, + } + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +class _FakeIndexConfigWithoutConcurrency: + """Stand-in for a PageIndex ``IndexConfig`` predating ``max_concurrency``.""" + + model_fields = { + "if_add_node_text": None, + "if_add_node_summary": None, + "if_add_doc_description": None, + } + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +class TestBuildIndexConfig: + def test_sets_base_flags(self): + cfg = _build_index_config({}) + assert cfg.if_add_node_text is True + assert cfg.if_add_node_summary is True + assert cfg.if_add_doc_description is True + + def test_forwards_concurrency_when_supported(self, monkeypatch): + monkeypatch.setattr("openkb.indexer.IndexConfig", _FakeIndexConfigWithConcurrency) + cfg = _build_index_config({"concurrency": 8}) + assert cfg.max_concurrency == 8 + + def test_does_not_forward_when_unsupported(self, monkeypatch): + monkeypatch.setattr("openkb.indexer.IndexConfig", _FakeIndexConfigWithoutConcurrency) + cfg = _build_index_config({"concurrency": 8}) + assert not hasattr(cfg, "max_concurrency") + + def test_none_value_is_left_to_pageindex_default(self, monkeypatch): + monkeypatch.setattr("openkb.indexer.IndexConfig", _FakeIndexConfigWithConcurrency) + cfg = _build_index_config({"concurrency": None}) + assert getattr(cfg, "max_concurrency", None) is None + + def test_invalid_value_is_left_to_pageindex_default(self, monkeypatch): + # resolve_concurrency() rejects bools/non-positive values — same as an + # unset key, just via the shared config-level validation. + monkeypatch.setattr("openkb.indexer.IndexConfig", _FakeIndexConfigWithConcurrency) + cfg = _build_index_config({"concurrency": 0}) + assert getattr(cfg, "max_concurrency", None) is None + + def test_warns_when_configured_but_unsupported(self, monkeypatch, caplog): + monkeypatch.setattr("openkb.indexer.IndexConfig", _FakeIndexConfigWithoutConcurrency) + with caplog.at_level(logging.WARNING, logger="openkb.indexer"): + _build_index_config({"concurrency": 8}) + assert "concurrency" in caplog.text + + def test_no_warning_when_unset(self, monkeypatch, caplog): + monkeypatch.setattr("openkb.indexer.IndexConfig", _FakeIndexConfigWithoutConcurrency) + with caplog.at_level(logging.WARNING, logger="openkb.indexer"): + _build_index_config({}) + assert caplog.text == "" + + def test_no_warning_when_supported(self, monkeypatch, caplog): + monkeypatch.setattr("openkb.indexer.IndexConfig", _FakeIndexConfigWithConcurrency) + with caplog.at_level(logging.WARNING, logger="openkb.indexer"): + _build_index_config({"concurrency": 8}) + assert caplog.text == "" class TestNormalizePageContent: @@ -201,6 +287,32 @@ def test_localclient_called_with_index_config(self, kb_dir, sample_tree, tmp_pat assert ic.if_add_node_summary is True assert ic.if_add_doc_description is True + def test_concurrency_flows_from_kb_config(self, kb_dir, sample_tree, tmp_path): + """The KB's real config.yaml, loaded by index_long_document itself, must + reach the IndexConfig passed to PageIndexClient — not just the isolated + _build_index_config unit tested directly with a hand-built dict.""" + (kb_dir / ".openkb" / "config.yaml").write_text( + "model: gpt-4o-mini\nconcurrency: 7\n", encoding="utf-8" + ) + + doc_id = "conc-789" + fake_col = self._make_fake_collection(doc_id, sample_tree) + fake_client = MagicMock() + fake_client.collection.return_value = fake_col + + pdf_path = tmp_path / "report.pdf" + pdf_path.write_bytes(b"%PDF-1.4 fake") + + with ( + patch("openkb.indexer.IndexConfig", _FakeIndexConfigWithConcurrency), + patch("openkb.indexer.PageIndexClient", 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) + + _, kwargs = mock_cls.call_args + assert kwargs["index_config"].max_concurrency == 7 + def test_cloud_page_content_is_normalized(self, kb_dir, sample_tree, tmp_path, monkeypatch): doc_id = "cloud-123" fake_col = self._make_fake_collection(doc_id, sample_tree) diff --git a/uv.lock b/uv.lock index bd4157c6..10f2a8b1 100644 --- a/uv.lock +++ b/uv.lock @@ -1975,7 +1975,7 @@ requires-dist = [ { name = "markitdown", extras = ["docx", "pptx", "xls", "xlsx"], specifier = "==0.1.5" }, { name = "mypy", marker = "extra == 'dev'", specifier = "==1.15.0" }, { name = "openai-agents", specifier = "==0.17.3" }, - { name = "pageindex", specifier = "==0.3.0.dev1" }, + { name = "pageindex", specifier = "==0.3.0.dev2" }, { name = "portalocker", specifier = "==3.2.0" }, { name = "prompt-toolkit", specifier = "==3.0.52" }, { name = "pytest", marker = "extra == 'dev'", specifier = "==9.0.3" }, @@ -2013,21 +2013,24 @@ wheels = [ [[package]] name = "pageindex" -version = "0.3.0.dev1" +version = "0.3.0.dev2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx", extra = ["socks"] }, { name = "litellm" }, + { name = "openai" }, { name = "openai-agents" }, + { name = "pydantic" }, { name = "pymupdf" }, { name = "pypdf2" }, { name = "python-dotenv" }, { name = "pyyaml" }, { name = "requests" }, + { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/53/27/318e73fe1ed2194b19abeb53f57f30b022c9ea4a3ebfe8c2b293f239d537/pageindex-0.3.0.dev1.tar.gz", hash = "sha256:77f02a85622a180918a11d4c4e4e97c10dc0ba142c6418ab1c501d6c1301a409", size = 61010, upload-time = "2026-04-10T17:20:58.524Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c2/87/9bf0e64e04509a453b110b09638eee297d8e046e8f709d13dc88d25a269b/pageindex-0.3.0.dev2.tar.gz", hash = "sha256:8d36899e5755710d82729f18bdb70f85de71cc2fcde0491749188ae6c8f3065f", size = 66284, upload-time = "2026-07-08T10:04:30.528Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/03/7b/a2d37a28475605b715f9fe28320989aec753c023b272fbf9c4bf9f693a55/pageindex-0.3.0.dev1-py3-none-any.whl", hash = "sha256:570c7fd7e54d62782030154dfe02b2f6b9e88407a293f96ea824203716ef28d3", size = 66954, upload-time = "2026-04-10T17:20:56.738Z" }, + { url = "https://files.pythonhosted.org/packages/93/99/7b24f5b7f36b6827f9c94d540248e79444103b5e2cb06ffe5ab0e8e8fe7f/pageindex-0.3.0.dev2-py3-none-any.whl", hash = "sha256:b0cb1f6e83e1a81d81f00b5b7bf229e1850fe54591fbb06d209fec0b0e63cf3b", size = 71889, upload-time = "2026-07-08T10:04:29.092Z" }, ] [[package]]