From 9180f4d91045ebad8e1696dd81aebccad7972c26 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Thu, 27 Aug 2026 15:27:26 +0200 Subject: [PATCH 01/16] fix(compiler): retry transient LLM API timeouts with bounded backoff - Added retryable exception classifier (_should_retry_exception) - Modified _llm_call() to pass retries=2 to litellm.completion() - Modified _llm_call_async() to pass retries=2 to litellm.acompletion() - LiteLLM handles exponential backoff (base 2) internally - Retries transient errors (Timeout, RateLimitError, ConnectionError) - Skips retry for permanent errors (ValueError, Auth, BadRequest) - Added comprehensive unit tests for exception filtering logic Fixes #229 --- openkb/agent/compiler.py | 128 ++++++++++++++++++++++++++++++++-- tests/test_compiler_retry.py | 130 +++++++++++++++++++++++++++++++++++ 2 files changed, 254 insertions(+), 4 deletions(-) create mode 100644 tests/test_compiler_retry.py diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878d..9ec19c32f 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -260,6 +260,71 @@ # --------------------------------------------------------------------------- +def _should_retry_exception(exc: Exception) -> bool: + """Determine whether an exception is retryable (transient error). + + Returns True for temporary API/network errors that may succeed on retry: + - Timeout (client-side or server-side) + - APIError 5xx (server errors) + - RateLimitError (429) + - ConnectionError / ServiceUnavailableError + + Returns False for permanent errors that won't be fixed by retry: + - TruncatedResponseError (model hit max_tokens) + - ValueError, TypeError (malformed input/output) + - AuthenticationError (credentials issue) + - BadRequestError (invalid parameters) + - Unknown error types (conservative approach) + """ + exc_type_name = type(exc).__name__ + + # ===== RETRYABLE (transient errors) ===== + + # Timeout (network/gateway timeout) + if "Timeout" in exc_type_name: + return True + + # Generic API errors (5xx range, but not 4xx) + if "APIError" in exc_type_name: + # Don't retry if it's a BadRequest/Invalid error (4xx) + if "Invalid" not in exc_type_name and "BadRequest" not in exc_type_name: + return True + + # Rate limiting (429) + if "RateLimitError" in exc_type_name or "Rate" in exc_type_name: + return True + + # Connection errors + if "ConnectionError" in exc_type_name: + return True + + # Service unavailable + if "ServiceUnavailable" in exc_type_name: + return True + + # ===== NOT RETRYABLE (permanent errors) ===== + + # Model hit max_tokens limit + if isinstance(exc, TruncatedResponseError): + return False + + # Content validation failures + if "ValueError" in exc_type_name or "TypeError" in exc_type_name: + return False + + # Authentication failures + if "Auth" in exc_type_name or "Permission" in exc_type_name: + return False + + # Bad parameters/requests + if "BadRequest" in exc_type_name or "Invalid" in exc_type_name: + return False + + # ===== UNKNOWN: Conservative approach ===== + # Don't retry errors we don't recognize + return False + + def _cached_text(text: str) -> list[dict]: """Wrap a text payload into a content-block list with an Anthropic ephemeral cache_control marker. @@ -406,7 +471,11 @@ def _llm_call( bundle=None, **kwargs, ) -> str: - """Single LLM call with animated progress and debug logging.""" + """Single LLM call with animated progress, debug logging, and retry support. + + Transient errors (Timeout, 5xx, 429) are automatically retried by LiteLLM. + Permanent errors (4xx, truncation, validation) are raised immediately. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -417,6 +486,10 @@ def _llm_call( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + + # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) + kwargs.setdefault("retries", 2) + logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -425,7 +498,27 @@ def _llm_call( spinner.start() t0 = time.time() - response = litellm.completion(model=model, messages=messages, **kwargs) + try: + response = litellm.completion(model=model, messages=messages, **kwargs) + except Exception as exc: + # NEW: Better error logging with retry context + if _should_retry_exception(exc): + logger.warning( + "LLM [%s] failed with transient error (retries applied by LiteLLM): %s", + step_name, + exc, + exc_info=False, # Don't spam stack traces for known transient errors + ) + else: + logger.warning( + "LLM [%s] failed with permanent error (no retry): %s", + step_name, + exc, + exc_info=True, # Full trace for unexpected errors + ) + spinner.stop("[FAILED]") + raise + content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) @@ -449,7 +542,11 @@ async def _llm_call_async( bundle=None, **kwargs, ) -> str: - """Async LLM call with timing output and debug logging.""" + """Async LLM call with timing output, debug logging, and retry support. + + Transient errors (Timeout, 5xx, 429) are automatically retried by LiteLLM. + Permanent errors (4xx, truncation, validation) are raised immediately. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -460,13 +557,36 @@ async def _llm_call_async( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + + # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) + kwargs.setdefault("retries", 2) + logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) t0 = time.time() - response = await litellm.acompletion(model=model, messages=messages, **kwargs) + try: + response = await litellm.acompletion(model=model, messages=messages, **kwargs) + except Exception as exc: + # NEW: Better error logging with retry context + if _should_retry_exception(exc): + logger.warning( + "LLM [%s] failed with transient error (retries applied by LiteLLM): %s", + step_name, + exc, + exc_info=False, + ) + else: + logger.warning( + "LLM [%s] failed with permanent error (no retry): %s", + step_name, + exc, + exc_info=True, + ) + raise + content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) diff --git a/tests/test_compiler_retry.py b/tests/test_compiler_retry.py new file mode 100644 index 000000000..7902201f7 --- /dev/null +++ b/tests/test_compiler_retry.py @@ -0,0 +1,130 @@ +"""Tests for LLM retry logic in compiler.py.""" + +from openkb.agent.compiler import TruncatedResponseError, _should_retry_exception + + +# Custom exception classes for testing (so we can control the type name) +class TimeoutError(Exception): + """Simulates litellm.Timeout.""" + + pass + + +class APIError(Exception): + """Simulates litellm.APIError (5xx).""" + + pass + + +class InvalidAPIError(APIError): + """Simulates InvalidAPIError (4xx).""" + + pass + + +class BadRequestError(Exception): + """Simulates BadRequestError.""" + + pass + + +class RateLimitError(Exception): + """Simulates litellm.RateLimitError.""" + + pass + + +class AuthenticationError(Exception): + """Simulates AuthenticationError.""" + + pass + + +class PermissionError(Exception): + """Simulates PermissionError.""" + + pass + + +class ServiceUnavailableError(Exception): + """Simulates ServiceUnavailableError.""" + + pass + + +class TestShouldRetryException: + """Test the exception filtering logic for retry decisions.""" + + def test_retryable_timeout(self): + """Timeout should be retryable.""" + exc = TimeoutError("Gateway Timeout") + assert _should_retry_exception(exc) is True + + def test_retryable_api_error_5xx(self): + """5xx API errors should be retryable.""" + exc = APIError("503 Service Unavailable") + assert _should_retry_exception(exc) is True + + def test_not_retryable_invalid_api_error(self): + """InvalidAPIError (4xx) should NOT be retryable.""" + exc = InvalidAPIError("400 Bad Request") + assert _should_retry_exception(exc) is False + + def test_retryable_rate_limit(self): + """Rate limit errors should be retryable.""" + exc = RateLimitError("429 Too Many Requests") + assert _should_retry_exception(exc) is True + + def test_retryable_connection_error(self): + """Connection errors should be retryable.""" + exc = ConnectionError("Connection refused") + assert _should_retry_exception(exc) is True + + def test_retryable_service_unavailable(self): + """Service unavailable errors should be retryable.""" + exc = ServiceUnavailableError("Service down") + assert _should_retry_exception(exc) is True + + def test_not_retryable_truncation(self): + """Truncated output should NOT be retryable.""" + exc = TruncatedResponseError("hit length limit") + assert _should_retry_exception(exc) is False + + def test_not_retryable_value_error(self): + """ValueError should NOT be retryable.""" + exc = ValueError("empty content") + assert _should_retry_exception(exc) is False + + def test_not_retryable_type_error(self): + """TypeError should NOT be retryable.""" + exc = TypeError("malformed") + assert _should_retry_exception(exc) is False + + def test_not_retryable_auth_error(self): + """Authentication errors should NOT be retryable.""" + exc = AuthenticationError("invalid API key") + assert _should_retry_exception(exc) is False + + def test_not_retryable_permission_error(self): + """Permission errors should NOT be retryable.""" + exc = PermissionError("forbidden") + assert _should_retry_exception(exc) is False + + def test_not_retryable_bad_request(self): + """BadRequest errors should NOT be retryable.""" + exc = BadRequestError("invalid params") + assert _should_retry_exception(exc) is False + + def test_not_retryable_unknown(self): + """Unknown errors should NOT be retried (conservative).""" + + class WeirdCustomError(Exception): + pass + + exc = WeirdCustomError("something weird") + assert _should_retry_exception(exc) is False + + def test_not_retryable_generic_exception(self): + """Generic Exception without special name should NOT be retried.""" + exc = Exception("generic error") + assert _should_retry_exception(exc) is False From cb78597a8f00de6af6324c7d902ae98d6afd5489 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Thu, 27 Aug 2026 15:53:37 +0200 Subject: [PATCH 02/16] feat(cli): add 'add-all' command and auto_delete_added_files config option - New 'openkb add-all' command processes all files in raw/ directory - New config parameter 'auto_delete_added_files' (default: false) - When enabled, both 'add' and 'add-all' automatically delete successfully ingested files - Updated help texts to document the new cleanup behavior - Config applies to all ingest methods: direct files, directories, and URLs --- openkb/cli.py | 90 ++++++++++++++++++++++++++++++++++++++++++++++-- openkb/config.py | 1 + 2 files changed, 89 insertions(+), 2 deletions(-) diff --git a/openkb/cli.py b/openkb/cli.py index c9e543181..1dcd81906 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -452,6 +452,29 @@ def add_single_file( return _add_single_file_locked(file_path, kb_dir, stage=stage, bundle=bundle) +def _delete_if_auto_cleanup_enabled( + file_path: Path, status: Literal["added", "skipped", "failed"], config: dict +) -> bool: + """Delete file if addition succeeded and auto_delete_added_files is enabled. + + Args: + file_path: Path to the file to potentially delete. + status: Result status from add_single_file ("added", "skipped", or "failed"). + config: Configuration dict (typically from resolve_effective_config). + + Returns: + True if file was deleted, False otherwise. + """ + if status == "added" and config.get("auto_delete_added_files", False): + try: + file_path.unlink(missing_ok=True) + return True + except Exception as exc: + logger.warning(f"Failed to delete {file_path.name}: {exc}") + return False + return False + + def _add_single_file_locked( file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None ) -> Literal["added", "skipped", "failed"]: @@ -1086,6 +1109,9 @@ def add(ctx, path, from_pageindex_cloud): Alternatively, pass --from-pageindex-cloud to import a document that is already indexed in PageIndex Cloud, with no local file. Requires the PAGEINDEX_API_KEY environment variable. + + If ``auto_delete_added_files`` is enabled in config.yaml, successfully + added files are automatically deleted after ingestion. """ kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) if kb_dir is None: @@ -1106,6 +1132,8 @@ def add(ctx, path, from_pageindex_cloud): click.echo("Provide a PATH or use --from-pageindex-cloud .") return + config = resolve_effective_config(kb_dir)[0] + # URL ingest: download into raw/ first, then call add_single_file explicitly. # Keep staged conversion enabled so converted source artifacts do not touch # the live KB before the mutation snapshot exists. The tri-state outcome @@ -1123,6 +1151,8 @@ def add(ctx, path, from_pageindex_cloud): # indexing has already succeeded but compilation didn't. if outcome == "skipped": fetched.unlink(missing_ok=True) + else: + _delete_if_auto_cleanup_enabled(fetched, outcome, config) return target = Path(path) @@ -1143,7 +1173,8 @@ def add(ctx, path, from_pageindex_cloud): click.echo(f"Found {total} supported file(s) in {path}.") for i, f in enumerate(files, 1): click.echo(f"\n[{i}/{total}] ", nl=False) - add_single_file(f, kb_dir) + outcome = add_single_file(f, kb_dir) + _delete_if_auto_cleanup_enabled(f, outcome, config) else: if target.suffix.lower() not in SUPPORTED_EXTENSIONS: click.echo( @@ -1151,7 +1182,62 @@ def add(ctx, path, from_pageindex_cloud): f"Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}" ) return - add_single_file(target, kb_dir) + outcome = add_single_file(target, kb_dir) + _delete_if_auto_cleanup_enabled(target, outcome, config) + + +@cli.command() +@click.pass_context +@_with_kb_lock(exclusive=True) +def add_all(ctx): + """Process all files in the ``raw/`` directory and add them to the knowledge base. + + This command walks the ``raw/`` directory recursively for all supported + document types and ingests them into the KB. If ``auto_delete_added_files`` + is enabled in config.yaml, successfully added files are automatically deleted + after ingestion. + + Returns a summary of the operation (added, skipped, failed, deleted counts). + """ + kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) + if kb_dir is None: + click.echo("No knowledge base found. Run `openkb init` first.") + return + + raw_dir = kb_dir / "raw" + if not raw_dir.is_dir(): + click.echo(f"No raw/ directory found at {raw_dir}") + return + + files = [ + f + for f in sorted(raw_dir.rglob("*")) + if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS + ] + if not files: + click.echo("No supported files found in raw/ directory.") + return + + config = resolve_effective_config(kb_dir)[0] + total = len(files) + added = skipped = failed = deleted = 0 + + click.echo(f"Processing {total} file(s) from raw/ directory...") + for i, f in enumerate(files, 1): + click.echo(f"\n[{i}/{total}] ", nl=False) + outcome = add_single_file(f, kb_dir) + if outcome == "added": + added += 1 + elif outcome == "skipped": + skipped += 1 + else: + failed += 1 + if _delete_if_auto_cleanup_enabled(f, outcome, config): + deleted += 1 + + click.echo( + f"\n\nSummary: Added: {added}, Skipped: {skipped}, Failed: {failed}, Deleted: {deleted}" + ) def _stream_to_tty() -> bool: diff --git a/openkb/config.py b/openkb/config.py index 95ca8691f..efd7ac382 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -36,6 +36,7 @@ # global/KB list overrides it wholesale; resolve_entity_types cleans the # effective value on read. "entity_types": list(DEFAULT_ENTITY_TYPES), + "auto_delete_added_files": False, } GLOBAL_CONFIG_DIR = Path.home() / ".config" / "openkb" From eecc0bc54e8cbe9c272f4ed787d29ea024c07344 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Thu, 27 Aug 2026 16:03:46 +0200 Subject: [PATCH 03/16] fix(cli): auto-delete also on 'skipped' status to keep raw/ clean Duplicates (skipped files) should also be auto-deleted when auto_delete_added_files is enabled, so raw/ stays clean. Only 'failed' status files are preserved to allow retries. Updated docstrings and helper function logic accordingly. --- openkb/cli.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/openkb/cli.py b/openkb/cli.py index 1dcd81906..0ade907e9 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -455,7 +455,10 @@ def add_single_file( def _delete_if_auto_cleanup_enabled( file_path: Path, status: Literal["added", "skipped", "failed"], config: dict ) -> bool: - """Delete file if addition succeeded and auto_delete_added_files is enabled. + """Delete file if auto_delete_added_files is enabled and ingestion succeeded/skipped. + + Deletes on both "added" (successful ingestion) and "skipped" (duplicate already + in KB) to keep raw/ directory clean. Preserves files on "failed" to allow retries. Args: file_path: Path to the file to potentially delete. @@ -465,7 +468,7 @@ def _delete_if_auto_cleanup_enabled( Returns: True if file was deleted, False otherwise. """ - if status == "added" and config.get("auto_delete_added_files", False): + if status in ("added", "skipped") and config.get("auto_delete_added_files", False): try: file_path.unlink(missing_ok=True) return True @@ -1110,8 +1113,9 @@ def add(ctx, path, from_pageindex_cloud): that is already indexed in PageIndex Cloud, with no local file. Requires the PAGEINDEX_API_KEY environment variable. - If ``auto_delete_added_files`` is enabled in config.yaml, successfully - added files are automatically deleted after ingestion. + If ``auto_delete_added_files`` is enabled in config.yaml, files are + automatically deleted after ingestion (both on successful addition and + on skip/duplicate). """ kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) if kb_dir is None: @@ -1194,8 +1198,8 @@ def add_all(ctx): This command walks the ``raw/`` directory recursively for all supported document types and ingests them into the KB. If ``auto_delete_added_files`` - is enabled in config.yaml, successfully added files are automatically deleted - after ingestion. + is enabled in config.yaml, files are automatically deleted after ingestion + (both on successful addition and on skip/duplicate). Returns a summary of the operation (added, skipped, failed, deleted counts). """ From 285d59a78c079468c75e13af4ff93a323ea14704 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Thu, 27 Aug 2026 18:06:06 +0200 Subject: [PATCH 04/16] fix(compiler): use LiteLLM's num_retries kwarg instead of retries The retry mechanism used kwargs.setdefault('retries', 2), but LiteLLM only recognizes 'num_retries'/'max_retries' as internal retry-control parameters. An unrecognized 'retries' kwarg falls through as a provider request-body field, which strict-mode proxies (e.g. custom Anthropic gateways) reject with 'retries: Extra inputs are not permitted'. Refines #229 --- openkb/agent/compiler.py | 18 ++++++++++--- tests/test_compiler_retry.py | 52 +++++++++++++++++++++++++++++++++++- 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index 9ec19c32f..2312f6902 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -487,8 +487,13 @@ def _llm_call( kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) - # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) - kwargs.setdefault("retries", 2) + # Retry configuration for transient errors (fixed: 2 retries). Uses + # LiteLLM's recognized ``num_retries`` kwarg — NOT ``retries``, which + # LiteLLM does not treat as an internal control parameter. An + # unrecognized kwarg falls through as a provider request-body field, + # which strict-mode proxies reject with e.g. "retries: Extra inputs + # are not permitted" (#233). + kwargs.setdefault("num_retries", 2) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: @@ -558,8 +563,13 @@ async def _llm_call_async( kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) - # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) - kwargs.setdefault("retries", 2) + # Retry configuration for transient errors (fixed: 2 retries). Uses + # LiteLLM's recognized ``num_retries`` kwarg — NOT ``retries``, which + # LiteLLM does not treat as an internal control parameter. An + # unrecognized kwarg falls through as a provider request-body field, + # which strict-mode proxies reject with e.g. "retries: Extra inputs + # are not permitted" (#233). + kwargs.setdefault("num_retries", 2) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: diff --git a/tests/test_compiler_retry.py b/tests/test_compiler_retry.py index 7902201f7..9ab7e3c69 100644 --- a/tests/test_compiler_retry.py +++ b/tests/test_compiler_retry.py @@ -1,6 +1,14 @@ """Tests for LLM retry logic in compiler.py.""" -from openkb.agent.compiler import TruncatedResponseError, _should_retry_exception +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +from openkb.agent.compiler import ( + TruncatedResponseError, + _llm_call, + _llm_call_async, + _should_retry_exception, +) # Custom exception classes for testing (so we can control the type name) @@ -128,3 +136,45 @@ def test_not_retryable_generic_exception(self): """Generic Exception without special name should NOT be retried.""" exc = Exception("generic error") assert _should_retry_exception(exc) is False + + +def _fake_response(): + choice = MagicMock() + choice.message.content = "ok" + choice.finish_reason = "stop" + resp = MagicMock() + resp.choices = [choice] + return resp + + +class TestRetryKwargForwarding: + """Regression tests for #233: the retry kwarg forwarded to LiteLLM must be + ``num_retries`` (LiteLLM's recognized internal control parameter), not + ``retries``. An unrecognized kwarg falls through as a provider + request-body field, which strict-mode proxies reject. + """ + + def test_llm_call_forwards_num_retries_not_retries(self): + with patch( + "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + ) as completion: + _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") + assert completion.call_args.kwargs["num_retries"] == 2 + assert "retries" not in completion.call_args.kwargs + + def test_llm_call_does_not_override_explicit_num_retries(self): + with patch( + "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + ) as completion: + _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step", num_retries=5) + assert completion.call_args.kwargs["num_retries"] == 5 + + def test_llm_call_async_forwards_num_retries_not_retries(self): + with patch( + "openkb.agent.compiler.litellm.acompletion", + new_callable=AsyncMock, + return_value=_fake_response(), + ) as acompletion: + asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step")) + assert acompletion.call_args.kwargs["num_retries"] == 2 + assert "retries" not in acompletion.call_args.kwargs From a3f970602a8530f54cc484ec25373f386e406aaf Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 15:38:19 +0200 Subject: [PATCH 05/16] fix(compiler): retry sweep + partial-compile status to prevent silent data loss - _compile_concepts now runs a deferred retry sweep after the first batch: concepts/entities that failed their first attempt (e.g. a transient Anthropic gateway timeout) are retried once more after the rest of the batch has run, giving real wall-clock time for the transient condition to clear while Anthropic's prompt cache is still warm. - compile_short_doc/compile_long_doc/_compile_concepts now return a bool reporting whether every planned concept/entity update was actually written. add_single_file/import_from_pageindex_cloud use this to add a new 'added_partial' status: the source hash is left unregistered and the raw file is kept (even with auto_delete_added_files) so a subsequent 'openkb add' retries the missing pieces instead of silently losing them. - Threaded the new status through the REST API (AddResponse.added_partial_count), the file watcher, and the add-all summary. Complements (does not replace) #229's num_retries fix: that fix corrects the immediate per-call retry mechanics; this adds an outer, delayed retry layer plus visibility/safety when even that is exhausted. --- openkb/agent/compiler.py | 119 +++++++++++++++++++++++++++++++----- openkb/api_helpers.py | 1 + openkb/api_models.py | 1 + openkb/cli.py | 125 +++++++++++++++++++++++++++----------- openkb/watch_service.py | 6 +- tests/test_add_command.py | 47 ++++++++++++++ tests/test_api.py | 1 + tests/test_compiler.py | 87 ++++++++++++++++++++++++++ 8 files changed, 337 insertions(+), 50 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index 2312f6902..dc8065968 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -26,6 +26,7 @@ import time import unicodedata from pathlib import Path +from typing import Any import litellm @@ -1734,7 +1735,7 @@ async def _compile_concepts( rewrite_summary: bool = False, entity_types: list[str] | None = None, bundle=None, -) -> None: +) -> bool: """Shared Steps 2-4: concepts plan → generate/update → index. Uses ``_CONCEPTS_PLAN_USER`` to get a plan with create/update/related @@ -1743,6 +1744,11 @@ async def _compile_concepts( written to disk. When ``rewrite_summary=True`` (short-doc path), the summary is rewritten by the LLM after concepts are finalized so its wikilinks reflect the actual concept pages on disk. + + Returns ``True`` when every planned concept and entity was written + (after the deferred retry sweep), ``False`` if any are still missing — + the caller uses this to decide whether the source document can be + considered fully compiled. """ source_file = f"summaries/{doc_name}.md" @@ -1819,7 +1825,9 @@ def _write_v1_summary_stripped() -> None: if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) - return + # Unparseable plan output means nothing was even attempted — treat as + # incomplete (not a "nothing to do" success) so the caller retries. + return False # Fallback: if LLM returns a flat list, treat all items as "create". # The new plan contract nests concepts under a "concepts" key alongside @@ -1839,7 +1847,9 @@ def _write_v1_summary_stripped() -> None: if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) - return + # Same as above: a scalar plan is a malformed-response case, not a + # deliberate "no concepts here" plan — treat as incomplete. + return False if isinstance(parsed, list): plan = {"create": _filter_concept_items(parsed, "list"), "update": [], "related": []} @@ -1921,7 +1931,12 @@ def _raw_group_count(group: object) -> int: if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) - return + # A genuinely empty plan (original_total == 0) is a complete, valid + # outcome — nothing was planned, so nothing is missing. But if items + # were planned and all got dropped as malformed (original_total > 0, + # already warned above), that's real content loss — report incomplete + # so the caller retries instead of silently accepting it. + return original_total == 0 # Build the whitelist of valid wikilink targets the LLM may emit. It # combines what already exists on disk with what *this* round will @@ -2093,17 +2108,30 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: _require_nonempty_content(content, name) return name, content, brief, etype_out - tasks = [] - tasks.extend(_gen_create(c) for c in create_items) - tasks.extend(_gen_update(c) for c in update_items) + # Kept alongside the coroutine lists below (in the same create-then-update + # order) so a failed item can be identified and retried by calling the + # same generator function again with its original plan dict — see the + # deferred retry sweep after the first gather. + concept_task_items: list[tuple[str, dict]] = [("create", c) for c in create_items] + [ + ("update", c) for c in update_items + ] + entity_task_items: list[tuple[str, dict]] = [("create", e) for e in entity_create] + [ + ("update", e) for e in entity_update + ] + + def _run_concept(kind: str, item: dict): + return _gen_create(item) if kind == "create" else _gen_update(item) + + def _run_entity(kind: str, item: dict): + return _gen_entity_create(item) if kind == "create" else _gen_entity_update(item) + + tasks = [_run_concept(kind, item) for kind, item in concept_task_items] # --- Step 3 (entities): build the entity task list up front so it can be # gathered concurrently with the concept tasks below. Entity coroutines # return 4-arity tuples (name, content, brief, type), so their results are # processed in their own loop rather than mixed with the concept tuples. - entity_tasks = [] - entity_tasks.extend(_gen_entity_create(e) for e in entity_create) - entity_tasks.extend(_gen_entity_update(e) for e in entity_update) + entity_tasks = [_run_entity(kind, item) for kind, item in entity_task_items] concept_names: list[str] = [] concept_briefs_map: dict[str, str] = {} @@ -2126,13 +2154,61 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: ) sys.stdout.flush() - results, entity_results = ([], []) + results: list[Any] = [] + entity_results: list[Any] = [] if tasks or entity_tasks: results, entity_results = await asyncio.gather( asyncio.gather(*tasks, return_exceptions=True), asyncio.gather(*entity_tasks, return_exceptions=True), ) + # --- Deferred retry sweep ------------------------------------------------- + # LiteLLM's own ``num_retries`` already retries a failing call, but those + # retries fire back-to-back against the very same transient provider hiccup + # (e.g. an Anthropic gateway timeout). Once every item in the batch has had + # its first attempt, retry only the ones that failed: by now real wall-clock + # time has passed (the rest of the batch ran in the meantime), giving a + # transient failure a real chance to have cleared, while Anthropic's prompt + # cache (~5 min TTL) is still warm from the sibling calls — so the retry is + # nearly as cheap as the original attempt. Bounded to a single extra sweep + # (not unbounded retries). + failed_concepts = [ + (kind, item) + for (kind, item), r in zip(concept_task_items, results) + if isinstance(r, Exception) + ] + failed_entities = [ + (kind, item) + for (kind, item), r in zip(entity_task_items, entity_results) + if isinstance(r, Exception) + ] + if failed_concepts or failed_entities: + sys.stdout.write( + f" Retrying {len(failed_concepts)} concept(s) and {len(failed_entities)} " + "entity(ies) that failed on the first attempt...\n" + ) + sys.stdout.flush() + retry_results, retry_entity_results = await asyncio.gather( + asyncio.gather( + *(_run_concept(kind, item) for kind, item in failed_concepts), + return_exceptions=True, + ), + asyncio.gather( + *(_run_entity(kind, item) for kind, item in failed_entities), + return_exceptions=True, + ), + ) + # Splice the retry outcomes back into the original result lists (in + # the same relative order the failures were collected in), so the + # processing below sees a single, already-reconciled result set and + # doesn't need to know a retry sweep happened. + retry_iter = iter(retry_results) + results = [next(retry_iter) if isinstance(r, Exception) else r for r in results] + retry_entity_iter = iter(retry_entity_results) + entity_results = [ + next(retry_entity_iter) if isinstance(r, Exception) else r for r in entity_results + ] + if tasks: failure_types: list[str] = [] for r in results: @@ -2331,6 +2407,13 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: entity_meta=entity_meta, ) + # True only if every planned concept and entity survived the retry sweep + # above. Callers (compile_short_doc/compile_long_doc → add_single_file) + # use this to avoid registering the source hash / deleting the raw file + # when some planned updates are still missing, so the next `openkb add` + # of the same file retries them instead of silently losing them. + return len(pending_writes) >= total and len(entity_pending) >= etotal + async def compile_short_doc( doc_name: str, @@ -2339,11 +2422,14 @@ async def compile_short_doc( model: str, max_concurrency: int = DEFAULT_COMPILE_CONCURRENCY, bundle=None, -) -> None: +) -> bool: """Compile a short document using a multi-step LLM pipeline with caching. Step 1: Build base context A (schema + doc content), generate summary. Steps 2-4: Delegated to ``_compile_concepts``. + + Returns ``True`` iff every planned concept/entity update was written + (see ``_compile_concepts``); ``False`` signals a partial compile. """ from openkb.config import resolve_effective_config @@ -2397,7 +2483,7 @@ async def compile_short_doc( # --- Steps 2-4: Concept plan → generate/update → summary rewrite → index --- try: - await _compile_concepts( + return await _compile_concepts( wiki_dir, kb_dir, model, @@ -2427,11 +2513,14 @@ async def compile_long_doc( doc_description: str = "", max_concurrency: int = DEFAULT_COMPILE_CONCURRENCY, bundle=None, -) -> None: +) -> bool: """Compile a long (PageIndex) document's concepts and index. The summary page is already written by the indexer. This function generates concept pages and updates the index. + + Returns ``True`` iff every planned concept/entity update was written + (see ``_compile_concepts``); ``False`` signals a partial compile. """ from openkb.config import resolve_effective_config @@ -2482,7 +2571,7 @@ async def compile_long_doc( # --- Steps 2-4: Concept plan → generate/update → index --- try: - await _compile_concepts( + return await _compile_concepts( wiki_dir, kb_dir, model, diff --git a/openkb/api_helpers.py b/openkb/api_helpers.py index d42bf4a9c..b1eb93b0d 100644 --- a/openkb/api_helpers.py +++ b/openkb/api_helpers.py @@ -277,6 +277,7 @@ def _summarize_add_results(kb: str, results: list[AddFileItem]) -> AddResponse: kb=kb, files=results, added_count=sum(1 for item in results if item.status == "added"), + added_partial_count=sum(1 for item in results if item.status == "added_partial"), skipped_count=sum(1 for item in results if item.status == "skipped"), failed_count=sum(1 for item in results if item.status == "failed"), ) diff --git a/openkb/api_models.py b/openkb/api_models.py index fb9d75755..f8c84ba28 100644 --- a/openkb/api_models.py +++ b/openkb/api_models.py @@ -118,6 +118,7 @@ class AddResponse(BaseModel): kb: str files: list[AddFileItem] added_count: int + added_partial_count: int = 0 skipped_count: int failed_count: int diff --git a/openkb/cli.py b/openkb/cli.py index 0ade907e9..7cca2c4fc 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -246,6 +246,13 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: "pageindex_cloud": "pageindex", } +# Outcome of a single-document add/import. "added_partial" means the document +# was ingested and every concept/entity that *did* generate was written, but +# one or more planned concept/entity updates were unrecoverable (even after +# the in-process retry sweep in ``_compile_concepts``) — the source hash is +# intentionally left unregistered so the file is retried, not silently lost. +AddStatus = Literal["added", "added_partial", "skipped", "failed"] + # Registry types that were compiled via the long-doc pipeline (tree + per-page # JSON source), as opposed to short docs (markdown source). Both the local # long-PDF type and cloud imports belong here — they share the long-doc @@ -428,12 +435,25 @@ def _snapshot_add_paths( return paths -def _run_compile_with_retry(coro_factory, label: str) -> None: +def _run_compile_with_retry(coro_factory, label: str) -> bool: + """Run a compile coroutine factory, retrying once on a hard exception. + + Returns the compile coroutine's own bool result (``True`` = every planned + concept/entity update was written, ``False`` = partial — see + ``compile_short_doc``/``compile_long_doc``). A ``None`` result (e.g. from + a test double that predates this contract) is treated as full success, + matching the old implicit "no exception raised = done" behavior. This is + orthogonal to the retry-on-exception loop below: a *raised* exception + here means the whole compile step blew up (plan call failed, etc.) and is + retried wholesale; a ``False`` return means the compile finished but some + individual concept/entity generations were unrecoverable even after their + own in-process retry sweep. + """ click.echo(f" {label}...") for attempt in range(2): try: - asyncio.run(coro_factory()) - return + result = asyncio.run(coro_factory()) + return True if result is None else bool(result) except Exception as exc: if attempt == 0: click.echo(" Retrying compilation in 2s...") @@ -442,27 +462,27 @@ def _run_compile_with_retry(coro_factory, label: str) -> None: click.echo(f" [ERROR] Compilation failed: {exc}") logger.debug("Compilation traceback:", exc_info=True) raise + return False # pragma: no cover - unreachable (loop always returns or raises) -def add_single_file( - file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None -) -> Literal["added", "skipped", "failed"]: +def add_single_file(file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None) -> AddStatus: """Convert, index, and compile a single document under the KB mutation lock.""" with kb_ingest_lock(kb_dir / ".openkb"): return _add_single_file_locked(file_path, kb_dir, stage=stage, bundle=bundle) -def _delete_if_auto_cleanup_enabled( - file_path: Path, status: Literal["added", "skipped", "failed"], config: dict -) -> bool: +def _delete_if_auto_cleanup_enabled(file_path: Path, status: AddStatus, config: dict) -> bool: """Delete file if auto_delete_added_files is enabled and ingestion succeeded/skipped. Deletes on both "added" (successful ingestion) and "skipped" (duplicate already - in KB) to keep raw/ directory clean. Preserves files on "failed" to allow retries. + in KB) to keep raw/ directory clean. Preserves files on "failed" — and on + "added_partial", where one or more planned concept/entity updates are still + missing — so the user (or the next `openkb add` run) can retry. Args: file_path: Path to the file to potentially delete. - status: Result status from add_single_file ("added", "skipped", or "failed"). + status: Result status from add_single_file ("added", "added_partial", + "skipped", or "failed"). config: Configuration dict (typically from resolve_effective_config). Returns: @@ -480,7 +500,7 @@ def _delete_if_auto_cleanup_enabled( def _add_single_file_locked( file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None -) -> Literal["added", "skipped", "failed"]: +) -> AddStatus: """Convert, index, and compile a single document into the knowledge base. Steps: @@ -490,12 +510,14 @@ def _add_single_file_locked( 4. Else: compile_short_doc. Returns: - ``"added"`` on full success, ``"skipped"`` when the file's hash - is already in the registry (dedup), or ``"failed"`` when any - pipeline stage raised. URL-ingest distinguishes these so it can - unlink the just-downloaded raw file on dedup (it would otherwise - be an orphan) while preserving it on failure so the user can - retry without re-downloading. + ``"added"`` on full success, ``"added_partial"`` when the document was + ingested but one or more planned concept/entity updates are still + missing (hash intentionally left unregistered so a re-add retries + them), ``"skipped"`` when the file's hash is already in the registry + (dedup), or ``"failed"`` when any pipeline stage raised. URL-ingest + distinguishes these so it can unlink the just-downloaded raw file on + dedup (it would otherwise be an orphan) while preserving it on + failure/partial success so the user can retry without re-downloading. """ from openkb.agent.compiler import compile_long_doc, compile_short_doc from openkb.state import HashRegistry @@ -527,11 +549,12 @@ def _add_single_file_locked( doc_name = result.doc_name or file_path.stem index_result = None # populated only on the long-doc branch + compile_ok = True # False if any planned concept/entity update is missing final_raw, final_source = _final_artifact_paths(result, kb_dir) def commit_body(snapshot) -> None: - nonlocal index_result + nonlocal index_result, compile_ok publish_staged_tree(staging_dir, kb_dir) if final_raw is not None: result.raw_path = final_raw @@ -576,7 +599,7 @@ def commit_body(snapshot) -> None: ) summary_path = kb_dir / "wiki" / "summaries" / f"{doc_name}.md" - _run_compile_with_retry( + compile_ok = _run_compile_with_retry( lambda: compile_long_doc( doc_name, summary_path, @@ -593,7 +616,7 @@ def commit_body(snapshot) -> None: if result.source_path is None: raise RuntimeError(f"Converted document has no source artifact: {file_path.name}") source_path = result.source_path - _run_compile_with_retry( + compile_ok = _run_compile_with_retry( lambda: compile_short_doc( doc_name, source_path, @@ -605,8 +628,12 @@ def commit_body(snapshot) -> None: label="Compiling short doc", ) - # Register hash only after successful compilation. - if result.file_hash: + # Register hash only after a fully successful compile. A partial + # compile (some planned concept/entity update still missing after + # the in-process retry sweep) intentionally leaves the hash + # unregistered so the next `openkb add` of this file is treated as + # new — not a duplicate to skip — and retries the missing pieces. + if compile_ok and result.file_hash: registry = HashRegistry(openkb_dir / "hashes.json") doc_type = "long_pdf" if result.is_long_doc else file_path.suffix.lstrip(".") meta = { @@ -654,8 +681,15 @@ def append_ingest_log() -> None: ) if not run_add_mutation(kb_dir, plan): return "failed" - click.echo(f" [OK] {file_path.name} added to knowledge base.") - return "added" + if compile_ok: + click.echo(f" [OK] {file_path.name} added to knowledge base.") + return "added" + click.echo( + f" [WARN] {file_path.name} added with some concept/entity updates missing " + "(see warnings above). Hash not registered — the next `openkb add` run will " + "retry the missing pieces." + ) + return "added_partial" @dataclass @@ -687,11 +721,16 @@ def _add_for_api(file_path: Path, kb_dir: Path, *, bundle=None) -> AddFileResult message = f"Already in knowledge base: {file_path.name}" elif status_str == "failed": message = f"Failed to add: {file_path.name} (see server logs)" + elif status_str == "added_partial": + message = ( + f"Added with some concept/entity updates missing: {file_path.name} " + "(will be retried on the next add)" + ) else: message = f"Added: {file_path.name}" return AddFileResult( original_name=file_path.name, - saved_path=str(file_path) if status_str == "added" else None, + saved_path=str(file_path) if status_str in ("added", "added_partial") else None, status=status_str, message=message, ) @@ -729,7 +768,7 @@ def _cleanup_failed_cloud_import(kb_dir: Path, doc_name: str) -> None: ) -def import_from_pageindex_cloud(doc_id: str, kb_dir: Path) -> Literal["added", "skipped", "failed"]: +def import_from_pageindex_cloud(doc_id: str, kb_dir: Path) -> AddStatus: """Import an existing PageIndex Cloud document into the KB by ``doc_id``. Fetches structure + page content from the cloud (no local PDF), compiles @@ -775,8 +814,10 @@ def import_from_pageindex_cloud(doc_id: str, kb_dir: Path) -> Literal["added", " stem = _cloud_display_stem(cloud.cloud_name, doc_id) doc_name = resolve_doc_name_from_key(stem, path_key, registry) + compile_ok = True # False if any planned concept/entity update is missing def commit_body(_snapshot) -> None: + nonlocal compile_ok summary_path = _write_long_doc_artifacts( cloud.tree, cloud.all_pages, @@ -785,7 +826,7 @@ def commit_body(_snapshot) -> None: kb_dir, description=cloud.description, ) - _run_compile_with_retry( + compile_ok = _run_compile_with_retry( lambda: compile_long_doc( doc_name, summary_path, @@ -798,7 +839,14 @@ def commit_body(_snapshot) -> None: label=f"Compiling imported doc (doc_id={doc_id})", ) - # Register the raw-less cloud entry only after successful compilation. + # Register the raw-less cloud entry only after a fully successful + # compile. There is no raw file fallback for a cloud import (the + # doc_id is deterministic), so registering on a partial compile + # would permanently mask the missing concept/entity updates — + # re-importing the same doc_id would just be treated as a dup + # and skipped forever. + if not compile_ok: + return registry = HashRegistry(openkb_dir / "hashes.json") meta = { "name": cloud.cloud_name, @@ -842,8 +890,14 @@ def append_cloud_log() -> None: logger.debug("Cloud import mutation traceback:", exc_info=True) return "failed" - click.echo(f" [OK] {doc_name} imported from PageIndex Cloud.") - return "added" + if compile_ok: + click.echo(f" [OK] {doc_name} imported from PageIndex Cloud.") + return "added" + click.echo( + f" [WARN] {doc_name} imported with some concept/entity updates missing " + "(see warnings above); hash not registered so re-running the import will retry them." + ) + return "added_partial" # --------------------------------------------------------------------------- @@ -1201,7 +1255,7 @@ def add_all(ctx): is enabled in config.yaml, files are automatically deleted after ingestion (both on successful addition and on skip/duplicate). - Returns a summary of the operation (added, skipped, failed, deleted counts). + Returns a summary of the operation (added, partial, skipped, failed, deleted counts). """ kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) if kb_dir is None: @@ -1224,7 +1278,7 @@ def add_all(ctx): config = resolve_effective_config(kb_dir)[0] total = len(files) - added = skipped = failed = deleted = 0 + added = partial = skipped = failed = deleted = 0 click.echo(f"Processing {total} file(s) from raw/ directory...") for i, f in enumerate(files, 1): @@ -1232,6 +1286,8 @@ def add_all(ctx): outcome = add_single_file(f, kb_dir) if outcome == "added": added += 1 + elif outcome == "added_partial": + partial += 1 elif outcome == "skipped": skipped += 1 else: @@ -1240,7 +1296,8 @@ def add_all(ctx): deleted += 1 click.echo( - f"\n\nSummary: Added: {added}, Skipped: {skipped}, Failed: {failed}, Deleted: {deleted}" + f"\n\nSummary: Added: {added}, Partial: {partial}, Skipped: {skipped}, " + f"Failed: {failed}, Deleted: {deleted}" ) diff --git a/openkb/watch_service.py b/openkb/watch_service.py index 1748a2004..24f43a598 100644 --- a/openkb/watch_service.py +++ b/openkb/watch_service.py @@ -143,7 +143,11 @@ def _process_file(state: WatcherState, raw_path: str) -> None: ) _inc(state, "failed") return - status = result.status if result.status in ("added", "skipped", "failed") else "failed" + status = ( + result.status + if result.status in ("added", "added_partial", "skipped", "failed") + else "failed" + ) _record_event( state, "file_done", diff --git a/tests/test_add_command.py b/tests/test_add_command.py index 3f51788e5..cb4321f16 100644 --- a/tests/test_add_command.py +++ b/tests/test_add_command.py @@ -372,6 +372,53 @@ async def compile_noop(*args, **kwargs): assert "path" in meta assert "stale-old-hash" not in hashes + def test_add_short_doc_partial_compile_skips_hash_and_keeps_file(self, tmp_path): + """When compile_short_doc reports an incomplete compile (some planned + concept/entity update was still missing after its own retry sweep), + the CLI must not register the hash or delete the raw file — so a + subsequent `openkb add` of the same file retries the missing pieces + instead of silently losing them.""" + kb_dir = self._setup_kb(tmp_path) + (kb_dir / ".openkb" / "config.yaml").write_text( + "model: gpt-4o-mini\nauto_delete_added_files: true\n" + ) + doc = tmp_path / "test.md" + doc.write_text("# Hello") + + source_path = kb_dir / "wiki" / "sources" / "test.md" + source_path.write_text("# Hello converted") + + from openkb.converter import ConvertResult + + mock_result = ConvertResult( + raw_path=kb_dir / "raw" / "test.md", + source_path=source_path, + is_long_doc=False, + file_hash="deadbeef00" * 8, + doc_name="test", + ) + + async def compile_partial(*args, **kwargs): + return False # simulates an unrecovered concept/entity failure + + runner = CliRunner() + with ( + patch("openkb.cli._find_kb_dir", return_value=kb_dir), + patch("openkb.cli.convert_document", return_value=mock_result), + patch("openkb.agent.compiler.compile_short_doc", new=compile_partial), + ): + result = runner.invoke(cli, ["add", str(doc)]) + + assert "WARN" in result.output + assert doc.exists(), "raw file must be kept (not auto-deleted) on a partial compile" + + import json as json_mod + + hashes = json_mod.loads((kb_dir / ".openkb" / "hashes.json").read_text(encoding="utf-8")) + assert mock_result.file_hash not in hashes, ( + "hash must stay unregistered on a partial compile so a re-add retries it" + ) + def test_add_oldest_legacy_entry_converges_to_single_entry(self, tmp_path): """Editing a pre-doc_name-era document must not fork the registry. diff --git a/tests/test_api.py b/tests/test_api.py index b9e0f91a2..1bdc25009 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -625,6 +625,7 @@ def fake_add(path, target_kb, **kwargs): }, ], "added_count": 1, + "added_partial_count": 0, "skipped_count": 1, "failed_count": 0, } diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 95a57cc4c..a76a35e57 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1882,6 +1882,93 @@ async def truncated_acompletion(*args, **kwargs): ) assert not (wiki / "concepts" / "ghost.md").exists(), "truncated create must be skipped" + @pytest.mark.asyncio + async def test_retry_sweep_recovers_concept_that_failed_first_attempt(self, tmp_path): + """A concept whose first attempt raises (simulating a transient + provider error, e.g. a gateway timeout) is retried once more after the + rest of the batch has run. If the retry succeeds, the page is written + and _compile_concepts reports full success (True).""" + wiki = self._setup_wiki(tmp_path) + plan_response = json.dumps( + { + "create": [ + {"name": "alpha", "title": "Alpha"}, + {"name": "beta", "title": "Beta"}, + ], + "update": [], + "related": [], + } + ) + alpha_response = json.dumps({"brief": "a", "content": "# Alpha\n\nRecovered on retry."}) + beta_response = json.dumps({"brief": "b", "content": "# Beta\n\nFine first try."}) + + call_count = {"n": 0} + + async def flaky_acompletion(*args, **kwargs): + idx = call_count["n"] + call_count["n"] += 1 + if idx == 0: + # First attempt overall (alpha, since max_concurrency=1 keeps + # this deterministic) fails like a transient gateway timeout. + raise RuntimeError("simulated transient gateway timeout") + mock_resp = MagicMock() + mock_resp.choices = [MagicMock()] + mock_resp.choices[0].message.content = beta_response if idx == 1 else alpha_response + mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) + mock_resp.usage.prompt_tokens_details = None + return mock_resp + + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) + mock_litellm.acompletion = AsyncMock(side_effect=flaky_acompletion) + result = await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 1, # max_concurrency=1 for deterministic call ordering + ) + + assert result is True, "the retry sweep should recover the failed concept" + assert (wiki / "concepts" / "beta.md").exists() + assert (wiki / "concepts" / "alpha.md").exists(), "alpha must be written after its retry" + assert "Recovered on retry." in (wiki / "concepts" / "alpha.md").read_text() + # alpha (fails) + beta (ok) + alpha retry (ok) = 3 acompletion calls. + assert call_count["n"] == 3 + + @pytest.mark.asyncio + async def test_compile_concepts_reports_partial_when_retry_also_fails(self, tmp_path): + """When a concept fails both its first attempt and the deferred retry, + _compile_concepts must report incomplete (False) instead of silently + treating the document as fully compiled.""" + wiki = self._setup_wiki(tmp_path) + plan_response = json.dumps( + {"create": [{"name": "ghost", "title": "Ghost"}], "update": [], "related": []} + ) + + async def always_fails(*args, **kwargs): + raise RuntimeError("simulated persistent gateway timeout") + + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) + mock_litellm.acompletion = AsyncMock(side_effect=always_fails) + result = await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 1, + ) + + assert result is False, "a concept that fails both attempts must be reported as incomplete" + assert not (wiki / "concepts" / "ghost.md").exists() + def test_page_fields_maps_response_shapes(self): """Shared mapping used by all four page closures: object, single-element array unwrap, wrong-shape skip (empty content), and non-JSON prose From 026b9c1584e3e170055820085934c9eebe5ef729 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 16:04:41 +0200 Subject: [PATCH 06/16] revert: back out retry-sweep + partial-status merge (fix/retry-sweep-and-partial-status) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Withdrawn pending root-cause investigation: the Gateway Timeout failures this and #229 target correlate with response duration/size (largest existing pages fail deterministically on every attempt), not random transient load — very likely an intermediary corporate LLM gateway/proxy enforcing a fixed request-duration limit, not something a client-side retry policy can fix. See related discussion on #229. The feature branch is kept (not deleted) on the fork for possible later reactivation once/if the actual root cause is addressed or the retry layer is judged worth keeping independently. --- openkb/agent/compiler.py | 119 +++++------------------------------- openkb/api_helpers.py | 1 - openkb/api_models.py | 1 - openkb/cli.py | 125 +++++++++++--------------------------- openkb/watch_service.py | 6 +- tests/test_add_command.py | 47 -------------- tests/test_api.py | 1 - tests/test_compiler.py | 87 -------------------------- 8 files changed, 50 insertions(+), 337 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index dc8065968..2312f6902 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -26,7 +26,6 @@ import time import unicodedata from pathlib import Path -from typing import Any import litellm @@ -1735,7 +1734,7 @@ async def _compile_concepts( rewrite_summary: bool = False, entity_types: list[str] | None = None, bundle=None, -) -> bool: +) -> None: """Shared Steps 2-4: concepts plan → generate/update → index. Uses ``_CONCEPTS_PLAN_USER`` to get a plan with create/update/related @@ -1744,11 +1743,6 @@ async def _compile_concepts( written to disk. When ``rewrite_summary=True`` (short-doc path), the summary is rewritten by the LLM after concepts are finalized so its wikilinks reflect the actual concept pages on disk. - - Returns ``True`` when every planned concept and entity was written - (after the deferred retry sweep), ``False`` if any are still missing — - the caller uses this to decide whether the source document can be - considered fully compiled. """ source_file = f"summaries/{doc_name}.md" @@ -1825,9 +1819,7 @@ def _write_v1_summary_stripped() -> None: if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) - # Unparseable plan output means nothing was even attempted — treat as - # incomplete (not a "nothing to do" success) so the caller retries. - return False + return # Fallback: if LLM returns a flat list, treat all items as "create". # The new plan contract nests concepts under a "concepts" key alongside @@ -1847,9 +1839,7 @@ def _write_v1_summary_stripped() -> None: if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) - # Same as above: a scalar plan is a malformed-response case, not a - # deliberate "no concepts here" plan — treat as incomplete. - return False + return if isinstance(parsed, list): plan = {"create": _filter_concept_items(parsed, "list"), "update": [], "related": []} @@ -1931,12 +1921,7 @@ def _raw_group_count(group: object) -> int: if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) - # A genuinely empty plan (original_total == 0) is a complete, valid - # outcome — nothing was planned, so nothing is missing. But if items - # were planned and all got dropped as malformed (original_total > 0, - # already warned above), that's real content loss — report incomplete - # so the caller retries instead of silently accepting it. - return original_total == 0 + return # Build the whitelist of valid wikilink targets the LLM may emit. It # combines what already exists on disk with what *this* round will @@ -2108,30 +2093,17 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: _require_nonempty_content(content, name) return name, content, brief, etype_out - # Kept alongside the coroutine lists below (in the same create-then-update - # order) so a failed item can be identified and retried by calling the - # same generator function again with its original plan dict — see the - # deferred retry sweep after the first gather. - concept_task_items: list[tuple[str, dict]] = [("create", c) for c in create_items] + [ - ("update", c) for c in update_items - ] - entity_task_items: list[tuple[str, dict]] = [("create", e) for e in entity_create] + [ - ("update", e) for e in entity_update - ] - - def _run_concept(kind: str, item: dict): - return _gen_create(item) if kind == "create" else _gen_update(item) - - def _run_entity(kind: str, item: dict): - return _gen_entity_create(item) if kind == "create" else _gen_entity_update(item) - - tasks = [_run_concept(kind, item) for kind, item in concept_task_items] + tasks = [] + tasks.extend(_gen_create(c) for c in create_items) + tasks.extend(_gen_update(c) for c in update_items) # --- Step 3 (entities): build the entity task list up front so it can be # gathered concurrently with the concept tasks below. Entity coroutines # return 4-arity tuples (name, content, brief, type), so their results are # processed in their own loop rather than mixed with the concept tuples. - entity_tasks = [_run_entity(kind, item) for kind, item in entity_task_items] + entity_tasks = [] + entity_tasks.extend(_gen_entity_create(e) for e in entity_create) + entity_tasks.extend(_gen_entity_update(e) for e in entity_update) concept_names: list[str] = [] concept_briefs_map: dict[str, str] = {} @@ -2154,61 +2126,13 @@ def _run_entity(kind: str, item: dict): ) sys.stdout.flush() - results: list[Any] = [] - entity_results: list[Any] = [] + results, entity_results = ([], []) if tasks or entity_tasks: results, entity_results = await asyncio.gather( asyncio.gather(*tasks, return_exceptions=True), asyncio.gather(*entity_tasks, return_exceptions=True), ) - # --- Deferred retry sweep ------------------------------------------------- - # LiteLLM's own ``num_retries`` already retries a failing call, but those - # retries fire back-to-back against the very same transient provider hiccup - # (e.g. an Anthropic gateway timeout). Once every item in the batch has had - # its first attempt, retry only the ones that failed: by now real wall-clock - # time has passed (the rest of the batch ran in the meantime), giving a - # transient failure a real chance to have cleared, while Anthropic's prompt - # cache (~5 min TTL) is still warm from the sibling calls — so the retry is - # nearly as cheap as the original attempt. Bounded to a single extra sweep - # (not unbounded retries). - failed_concepts = [ - (kind, item) - for (kind, item), r in zip(concept_task_items, results) - if isinstance(r, Exception) - ] - failed_entities = [ - (kind, item) - for (kind, item), r in zip(entity_task_items, entity_results) - if isinstance(r, Exception) - ] - if failed_concepts or failed_entities: - sys.stdout.write( - f" Retrying {len(failed_concepts)} concept(s) and {len(failed_entities)} " - "entity(ies) that failed on the first attempt...\n" - ) - sys.stdout.flush() - retry_results, retry_entity_results = await asyncio.gather( - asyncio.gather( - *(_run_concept(kind, item) for kind, item in failed_concepts), - return_exceptions=True, - ), - asyncio.gather( - *(_run_entity(kind, item) for kind, item in failed_entities), - return_exceptions=True, - ), - ) - # Splice the retry outcomes back into the original result lists (in - # the same relative order the failures were collected in), so the - # processing below sees a single, already-reconciled result set and - # doesn't need to know a retry sweep happened. - retry_iter = iter(retry_results) - results = [next(retry_iter) if isinstance(r, Exception) else r for r in results] - retry_entity_iter = iter(retry_entity_results) - entity_results = [ - next(retry_entity_iter) if isinstance(r, Exception) else r for r in entity_results - ] - if tasks: failure_types: list[str] = [] for r in results: @@ -2407,13 +2331,6 @@ def _run_entity(kind: str, item: dict): entity_meta=entity_meta, ) - # True only if every planned concept and entity survived the retry sweep - # above. Callers (compile_short_doc/compile_long_doc → add_single_file) - # use this to avoid registering the source hash / deleting the raw file - # when some planned updates are still missing, so the next `openkb add` - # of the same file retries them instead of silently losing them. - return len(pending_writes) >= total and len(entity_pending) >= etotal - async def compile_short_doc( doc_name: str, @@ -2422,14 +2339,11 @@ async def compile_short_doc( model: str, max_concurrency: int = DEFAULT_COMPILE_CONCURRENCY, bundle=None, -) -> bool: +) -> None: """Compile a short document using a multi-step LLM pipeline with caching. Step 1: Build base context A (schema + doc content), generate summary. Steps 2-4: Delegated to ``_compile_concepts``. - - Returns ``True`` iff every planned concept/entity update was written - (see ``_compile_concepts``); ``False`` signals a partial compile. """ from openkb.config import resolve_effective_config @@ -2483,7 +2397,7 @@ async def compile_short_doc( # --- Steps 2-4: Concept plan → generate/update → summary rewrite → index --- try: - return await _compile_concepts( + await _compile_concepts( wiki_dir, kb_dir, model, @@ -2513,14 +2427,11 @@ async def compile_long_doc( doc_description: str = "", max_concurrency: int = DEFAULT_COMPILE_CONCURRENCY, bundle=None, -) -> bool: +) -> None: """Compile a long (PageIndex) document's concepts and index. The summary page is already written by the indexer. This function generates concept pages and updates the index. - - Returns ``True`` iff every planned concept/entity update was written - (see ``_compile_concepts``); ``False`` signals a partial compile. """ from openkb.config import resolve_effective_config @@ -2571,7 +2482,7 @@ async def compile_long_doc( # --- Steps 2-4: Concept plan → generate/update → index --- try: - return await _compile_concepts( + await _compile_concepts( wiki_dir, kb_dir, model, diff --git a/openkb/api_helpers.py b/openkb/api_helpers.py index b1eb93b0d..d42bf4a9c 100644 --- a/openkb/api_helpers.py +++ b/openkb/api_helpers.py @@ -277,7 +277,6 @@ def _summarize_add_results(kb: str, results: list[AddFileItem]) -> AddResponse: kb=kb, files=results, added_count=sum(1 for item in results if item.status == "added"), - added_partial_count=sum(1 for item in results if item.status == "added_partial"), skipped_count=sum(1 for item in results if item.status == "skipped"), failed_count=sum(1 for item in results if item.status == "failed"), ) diff --git a/openkb/api_models.py b/openkb/api_models.py index f8c84ba28..fb9d75755 100644 --- a/openkb/api_models.py +++ b/openkb/api_models.py @@ -118,7 +118,6 @@ class AddResponse(BaseModel): kb: str files: list[AddFileItem] added_count: int - added_partial_count: int = 0 skipped_count: int failed_count: int diff --git a/openkb/cli.py b/openkb/cli.py index 7cca2c4fc..0ade907e9 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -246,13 +246,6 @@ def _setup_llm_key(kb_dir: Path | None = None) -> None: "pageindex_cloud": "pageindex", } -# Outcome of a single-document add/import. "added_partial" means the document -# was ingested and every concept/entity that *did* generate was written, but -# one or more planned concept/entity updates were unrecoverable (even after -# the in-process retry sweep in ``_compile_concepts``) — the source hash is -# intentionally left unregistered so the file is retried, not silently lost. -AddStatus = Literal["added", "added_partial", "skipped", "failed"] - # Registry types that were compiled via the long-doc pipeline (tree + per-page # JSON source), as opposed to short docs (markdown source). Both the local # long-PDF type and cloud imports belong here — they share the long-doc @@ -435,25 +428,12 @@ def _snapshot_add_paths( return paths -def _run_compile_with_retry(coro_factory, label: str) -> bool: - """Run a compile coroutine factory, retrying once on a hard exception. - - Returns the compile coroutine's own bool result (``True`` = every planned - concept/entity update was written, ``False`` = partial — see - ``compile_short_doc``/``compile_long_doc``). A ``None`` result (e.g. from - a test double that predates this contract) is treated as full success, - matching the old implicit "no exception raised = done" behavior. This is - orthogonal to the retry-on-exception loop below: a *raised* exception - here means the whole compile step blew up (plan call failed, etc.) and is - retried wholesale; a ``False`` return means the compile finished but some - individual concept/entity generations were unrecoverable even after their - own in-process retry sweep. - """ +def _run_compile_with_retry(coro_factory, label: str) -> None: click.echo(f" {label}...") for attempt in range(2): try: - result = asyncio.run(coro_factory()) - return True if result is None else bool(result) + asyncio.run(coro_factory()) + return except Exception as exc: if attempt == 0: click.echo(" Retrying compilation in 2s...") @@ -462,27 +442,27 @@ def _run_compile_with_retry(coro_factory, label: str) -> bool: click.echo(f" [ERROR] Compilation failed: {exc}") logger.debug("Compilation traceback:", exc_info=True) raise - return False # pragma: no cover - unreachable (loop always returns or raises) -def add_single_file(file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None) -> AddStatus: +def add_single_file( + file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None +) -> Literal["added", "skipped", "failed"]: """Convert, index, and compile a single document under the KB mutation lock.""" with kb_ingest_lock(kb_dir / ".openkb"): return _add_single_file_locked(file_path, kb_dir, stage=stage, bundle=bundle) -def _delete_if_auto_cleanup_enabled(file_path: Path, status: AddStatus, config: dict) -> bool: +def _delete_if_auto_cleanup_enabled( + file_path: Path, status: Literal["added", "skipped", "failed"], config: dict +) -> bool: """Delete file if auto_delete_added_files is enabled and ingestion succeeded/skipped. Deletes on both "added" (successful ingestion) and "skipped" (duplicate already - in KB) to keep raw/ directory clean. Preserves files on "failed" — and on - "added_partial", where one or more planned concept/entity updates are still - missing — so the user (or the next `openkb add` run) can retry. + in KB) to keep raw/ directory clean. Preserves files on "failed" to allow retries. Args: file_path: Path to the file to potentially delete. - status: Result status from add_single_file ("added", "added_partial", - "skipped", or "failed"). + status: Result status from add_single_file ("added", "skipped", or "failed"). config: Configuration dict (typically from resolve_effective_config). Returns: @@ -500,7 +480,7 @@ def _delete_if_auto_cleanup_enabled(file_path: Path, status: AddStatus, config: def _add_single_file_locked( file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None -) -> AddStatus: +) -> Literal["added", "skipped", "failed"]: """Convert, index, and compile a single document into the knowledge base. Steps: @@ -510,14 +490,12 @@ def _add_single_file_locked( 4. Else: compile_short_doc. Returns: - ``"added"`` on full success, ``"added_partial"`` when the document was - ingested but one or more planned concept/entity updates are still - missing (hash intentionally left unregistered so a re-add retries - them), ``"skipped"`` when the file's hash is already in the registry - (dedup), or ``"failed"`` when any pipeline stage raised. URL-ingest - distinguishes these so it can unlink the just-downloaded raw file on - dedup (it would otherwise be an orphan) while preserving it on - failure/partial success so the user can retry without re-downloading. + ``"added"`` on full success, ``"skipped"`` when the file's hash + is already in the registry (dedup), or ``"failed"`` when any + pipeline stage raised. URL-ingest distinguishes these so it can + unlink the just-downloaded raw file on dedup (it would otherwise + be an orphan) while preserving it on failure so the user can + retry without re-downloading. """ from openkb.agent.compiler import compile_long_doc, compile_short_doc from openkb.state import HashRegistry @@ -549,12 +527,11 @@ def _add_single_file_locked( doc_name = result.doc_name or file_path.stem index_result = None # populated only on the long-doc branch - compile_ok = True # False if any planned concept/entity update is missing final_raw, final_source = _final_artifact_paths(result, kb_dir) def commit_body(snapshot) -> None: - nonlocal index_result, compile_ok + nonlocal index_result publish_staged_tree(staging_dir, kb_dir) if final_raw is not None: result.raw_path = final_raw @@ -599,7 +576,7 @@ def commit_body(snapshot) -> None: ) summary_path = kb_dir / "wiki" / "summaries" / f"{doc_name}.md" - compile_ok = _run_compile_with_retry( + _run_compile_with_retry( lambda: compile_long_doc( doc_name, summary_path, @@ -616,7 +593,7 @@ def commit_body(snapshot) -> None: if result.source_path is None: raise RuntimeError(f"Converted document has no source artifact: {file_path.name}") source_path = result.source_path - compile_ok = _run_compile_with_retry( + _run_compile_with_retry( lambda: compile_short_doc( doc_name, source_path, @@ -628,12 +605,8 @@ def commit_body(snapshot) -> None: label="Compiling short doc", ) - # Register hash only after a fully successful compile. A partial - # compile (some planned concept/entity update still missing after - # the in-process retry sweep) intentionally leaves the hash - # unregistered so the next `openkb add` of this file is treated as - # new — not a duplicate to skip — and retries the missing pieces. - if compile_ok and result.file_hash: + # Register hash only after successful compilation. + if result.file_hash: registry = HashRegistry(openkb_dir / "hashes.json") doc_type = "long_pdf" if result.is_long_doc else file_path.suffix.lstrip(".") meta = { @@ -681,15 +654,8 @@ def append_ingest_log() -> None: ) if not run_add_mutation(kb_dir, plan): return "failed" - if compile_ok: - click.echo(f" [OK] {file_path.name} added to knowledge base.") - return "added" - click.echo( - f" [WARN] {file_path.name} added with some concept/entity updates missing " - "(see warnings above). Hash not registered — the next `openkb add` run will " - "retry the missing pieces." - ) - return "added_partial" + click.echo(f" [OK] {file_path.name} added to knowledge base.") + return "added" @dataclass @@ -721,16 +687,11 @@ def _add_for_api(file_path: Path, kb_dir: Path, *, bundle=None) -> AddFileResult message = f"Already in knowledge base: {file_path.name}" elif status_str == "failed": message = f"Failed to add: {file_path.name} (see server logs)" - elif status_str == "added_partial": - message = ( - f"Added with some concept/entity updates missing: {file_path.name} " - "(will be retried on the next add)" - ) else: message = f"Added: {file_path.name}" return AddFileResult( original_name=file_path.name, - saved_path=str(file_path) if status_str in ("added", "added_partial") else None, + saved_path=str(file_path) if status_str == "added" else None, status=status_str, message=message, ) @@ -768,7 +729,7 @@ def _cleanup_failed_cloud_import(kb_dir: Path, doc_name: str) -> None: ) -def import_from_pageindex_cloud(doc_id: str, kb_dir: Path) -> AddStatus: +def import_from_pageindex_cloud(doc_id: str, kb_dir: Path) -> Literal["added", "skipped", "failed"]: """Import an existing PageIndex Cloud document into the KB by ``doc_id``. Fetches structure + page content from the cloud (no local PDF), compiles @@ -814,10 +775,8 @@ def import_from_pageindex_cloud(doc_id: str, kb_dir: Path) -> AddStatus: stem = _cloud_display_stem(cloud.cloud_name, doc_id) doc_name = resolve_doc_name_from_key(stem, path_key, registry) - compile_ok = True # False if any planned concept/entity update is missing def commit_body(_snapshot) -> None: - nonlocal compile_ok summary_path = _write_long_doc_artifacts( cloud.tree, cloud.all_pages, @@ -826,7 +785,7 @@ def commit_body(_snapshot) -> None: kb_dir, description=cloud.description, ) - compile_ok = _run_compile_with_retry( + _run_compile_with_retry( lambda: compile_long_doc( doc_name, summary_path, @@ -839,14 +798,7 @@ def commit_body(_snapshot) -> None: label=f"Compiling imported doc (doc_id={doc_id})", ) - # Register the raw-less cloud entry only after a fully successful - # compile. There is no raw file fallback for a cloud import (the - # doc_id is deterministic), so registering on a partial compile - # would permanently mask the missing concept/entity updates — - # re-importing the same doc_id would just be treated as a dup - # and skipped forever. - if not compile_ok: - return + # Register the raw-less cloud entry only after successful compilation. registry = HashRegistry(openkb_dir / "hashes.json") meta = { "name": cloud.cloud_name, @@ -890,14 +842,8 @@ def append_cloud_log() -> None: logger.debug("Cloud import mutation traceback:", exc_info=True) return "failed" - if compile_ok: - click.echo(f" [OK] {doc_name} imported from PageIndex Cloud.") - return "added" - click.echo( - f" [WARN] {doc_name} imported with some concept/entity updates missing " - "(see warnings above); hash not registered so re-running the import will retry them." - ) - return "added_partial" + click.echo(f" [OK] {doc_name} imported from PageIndex Cloud.") + return "added" # --------------------------------------------------------------------------- @@ -1255,7 +1201,7 @@ def add_all(ctx): is enabled in config.yaml, files are automatically deleted after ingestion (both on successful addition and on skip/duplicate). - Returns a summary of the operation (added, partial, skipped, failed, deleted counts). + Returns a summary of the operation (added, skipped, failed, deleted counts). """ kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) if kb_dir is None: @@ -1278,7 +1224,7 @@ def add_all(ctx): config = resolve_effective_config(kb_dir)[0] total = len(files) - added = partial = skipped = failed = deleted = 0 + added = skipped = failed = deleted = 0 click.echo(f"Processing {total} file(s) from raw/ directory...") for i, f in enumerate(files, 1): @@ -1286,8 +1232,6 @@ def add_all(ctx): outcome = add_single_file(f, kb_dir) if outcome == "added": added += 1 - elif outcome == "added_partial": - partial += 1 elif outcome == "skipped": skipped += 1 else: @@ -1296,8 +1240,7 @@ def add_all(ctx): deleted += 1 click.echo( - f"\n\nSummary: Added: {added}, Partial: {partial}, Skipped: {skipped}, " - f"Failed: {failed}, Deleted: {deleted}" + f"\n\nSummary: Added: {added}, Skipped: {skipped}, Failed: {failed}, Deleted: {deleted}" ) diff --git a/openkb/watch_service.py b/openkb/watch_service.py index 24f43a598..1748a2004 100644 --- a/openkb/watch_service.py +++ b/openkb/watch_service.py @@ -143,11 +143,7 @@ def _process_file(state: WatcherState, raw_path: str) -> None: ) _inc(state, "failed") return - status = ( - result.status - if result.status in ("added", "added_partial", "skipped", "failed") - else "failed" - ) + status = result.status if result.status in ("added", "skipped", "failed") else "failed" _record_event( state, "file_done", diff --git a/tests/test_add_command.py b/tests/test_add_command.py index cb4321f16..3f51788e5 100644 --- a/tests/test_add_command.py +++ b/tests/test_add_command.py @@ -372,53 +372,6 @@ async def compile_noop(*args, **kwargs): assert "path" in meta assert "stale-old-hash" not in hashes - def test_add_short_doc_partial_compile_skips_hash_and_keeps_file(self, tmp_path): - """When compile_short_doc reports an incomplete compile (some planned - concept/entity update was still missing after its own retry sweep), - the CLI must not register the hash or delete the raw file — so a - subsequent `openkb add` of the same file retries the missing pieces - instead of silently losing them.""" - kb_dir = self._setup_kb(tmp_path) - (kb_dir / ".openkb" / "config.yaml").write_text( - "model: gpt-4o-mini\nauto_delete_added_files: true\n" - ) - doc = tmp_path / "test.md" - doc.write_text("# Hello") - - source_path = kb_dir / "wiki" / "sources" / "test.md" - source_path.write_text("# Hello converted") - - from openkb.converter import ConvertResult - - mock_result = ConvertResult( - raw_path=kb_dir / "raw" / "test.md", - source_path=source_path, - is_long_doc=False, - file_hash="deadbeef00" * 8, - doc_name="test", - ) - - async def compile_partial(*args, **kwargs): - return False # simulates an unrecovered concept/entity failure - - runner = CliRunner() - with ( - patch("openkb.cli._find_kb_dir", return_value=kb_dir), - patch("openkb.cli.convert_document", return_value=mock_result), - patch("openkb.agent.compiler.compile_short_doc", new=compile_partial), - ): - result = runner.invoke(cli, ["add", str(doc)]) - - assert "WARN" in result.output - assert doc.exists(), "raw file must be kept (not auto-deleted) on a partial compile" - - import json as json_mod - - hashes = json_mod.loads((kb_dir / ".openkb" / "hashes.json").read_text(encoding="utf-8")) - assert mock_result.file_hash not in hashes, ( - "hash must stay unregistered on a partial compile so a re-add retries it" - ) - def test_add_oldest_legacy_entry_converges_to_single_entry(self, tmp_path): """Editing a pre-doc_name-era document must not fork the registry. diff --git a/tests/test_api.py b/tests/test_api.py index 1bdc25009..b9e0f91a2 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -625,7 +625,6 @@ def fake_add(path, target_kb, **kwargs): }, ], "added_count": 1, - "added_partial_count": 0, "skipped_count": 1, "failed_count": 0, } diff --git a/tests/test_compiler.py b/tests/test_compiler.py index a76a35e57..95a57cc4c 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1882,93 +1882,6 @@ async def truncated_acompletion(*args, **kwargs): ) assert not (wiki / "concepts" / "ghost.md").exists(), "truncated create must be skipped" - @pytest.mark.asyncio - async def test_retry_sweep_recovers_concept_that_failed_first_attempt(self, tmp_path): - """A concept whose first attempt raises (simulating a transient - provider error, e.g. a gateway timeout) is retried once more after the - rest of the batch has run. If the retry succeeds, the page is written - and _compile_concepts reports full success (True).""" - wiki = self._setup_wiki(tmp_path) - plan_response = json.dumps( - { - "create": [ - {"name": "alpha", "title": "Alpha"}, - {"name": "beta", "title": "Beta"}, - ], - "update": [], - "related": [], - } - ) - alpha_response = json.dumps({"brief": "a", "content": "# Alpha\n\nRecovered on retry."}) - beta_response = json.dumps({"brief": "b", "content": "# Beta\n\nFine first try."}) - - call_count = {"n": 0} - - async def flaky_acompletion(*args, **kwargs): - idx = call_count["n"] - call_count["n"] += 1 - if idx == 0: - # First attempt overall (alpha, since max_concurrency=1 keeps - # this deterministic) fails like a transient gateway timeout. - raise RuntimeError("simulated transient gateway timeout") - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = beta_response if idx == 1 else alpha_response - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp - - with patch("openkb.agent.compiler.litellm") as mock_litellm: - mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) - mock_litellm.acompletion = AsyncMock(side_effect=flaky_acompletion) - result = await _compile_concepts( - wiki, - tmp_path, - "gpt-4o-mini", - {"role": "system", "content": "s"}, - {"role": "user", "content": "d"}, - "summary", - "test-doc", - 1, # max_concurrency=1 for deterministic call ordering - ) - - assert result is True, "the retry sweep should recover the failed concept" - assert (wiki / "concepts" / "beta.md").exists() - assert (wiki / "concepts" / "alpha.md").exists(), "alpha must be written after its retry" - assert "Recovered on retry." in (wiki / "concepts" / "alpha.md").read_text() - # alpha (fails) + beta (ok) + alpha retry (ok) = 3 acompletion calls. - assert call_count["n"] == 3 - - @pytest.mark.asyncio - async def test_compile_concepts_reports_partial_when_retry_also_fails(self, tmp_path): - """When a concept fails both its first attempt and the deferred retry, - _compile_concepts must report incomplete (False) instead of silently - treating the document as fully compiled.""" - wiki = self._setup_wiki(tmp_path) - plan_response = json.dumps( - {"create": [{"name": "ghost", "title": "Ghost"}], "update": [], "related": []} - ) - - async def always_fails(*args, **kwargs): - raise RuntimeError("simulated persistent gateway timeout") - - with patch("openkb.agent.compiler.litellm") as mock_litellm: - mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) - mock_litellm.acompletion = AsyncMock(side_effect=always_fails) - result = await _compile_concepts( - wiki, - tmp_path, - "gpt-4o-mini", - {"role": "system", "content": "s"}, - {"role": "user", "content": "d"}, - "summary", - "test-doc", - 1, - ) - - assert result is False, "a concept that fails both attempts must be reported as incomplete" - assert not (wiki / "concepts" / "ghost.md").exists() - def test_page_fields_maps_response_shapes(self): """Shared mapping used by all four page closures: object, single-element array unwrap, wrong-shape skip (empty content), and non-JSON prose From f8a546679da7902389371bd69eb3bab2246b6279 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 16:04:50 +0200 Subject: [PATCH 07/16] revert: back out num_retries fix merge (fix/issue-229-retry-timeout, #229) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Withdrawn together with the retry-sweep work (see prior revert commit): root-cause investigation shows the Gateway Timeout failures correlate with response duration/size (the largest existing pages fail deterministically on every attempt regardless of retries), pointing to an intermediary corporate LLM gateway/proxy enforcing a fixed request-duration limit rather than transient upstream load. A client-side retry policy does not address that root cause, and other users behind a similar corporate proxy would likely hit the same wall first — so retrying isn't the right fix to ship right now. The fix itself (num_retries vs. the unrecognized retries kwarg) may still be independently correct/useful; the feature branch is kept (not deleted) on the fork for possible later reactivation. Related upstream: closes our own copies tracking VectifyAI/OpenKB#229 and VectifyAI/OpenKB#230 (withdrawn there directly, not merged). --- openkb/agent/compiler.py | 18 +++---------- tests/test_compiler_retry.py | 52 +----------------------------------- 2 files changed, 5 insertions(+), 65 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index 2312f6902..9ec19c32f 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -487,13 +487,8 @@ def _llm_call( kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) - # Retry configuration for transient errors (fixed: 2 retries). Uses - # LiteLLM's recognized ``num_retries`` kwarg — NOT ``retries``, which - # LiteLLM does not treat as an internal control parameter. An - # unrecognized kwarg falls through as a provider request-body field, - # which strict-mode proxies reject with e.g. "retries: Extra inputs - # are not permitted" (#233). - kwargs.setdefault("num_retries", 2) + # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) + kwargs.setdefault("retries", 2) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: @@ -563,13 +558,8 @@ async def _llm_call_async( kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) - # Retry configuration for transient errors (fixed: 2 retries). Uses - # LiteLLM's recognized ``num_retries`` kwarg — NOT ``retries``, which - # LiteLLM does not treat as an internal control parameter. An - # unrecognized kwarg falls through as a provider request-body field, - # which strict-mode proxies reject with e.g. "retries: Extra inputs - # are not permitted" (#233). - kwargs.setdefault("num_retries", 2) + # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) + kwargs.setdefault("retries", 2) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: diff --git a/tests/test_compiler_retry.py b/tests/test_compiler_retry.py index 9ab7e3c69..7902201f7 100644 --- a/tests/test_compiler_retry.py +++ b/tests/test_compiler_retry.py @@ -1,14 +1,6 @@ """Tests for LLM retry logic in compiler.py.""" -import asyncio -from unittest.mock import AsyncMock, MagicMock, patch - -from openkb.agent.compiler import ( - TruncatedResponseError, - _llm_call, - _llm_call_async, - _should_retry_exception, -) +from openkb.agent.compiler import TruncatedResponseError, _should_retry_exception # Custom exception classes for testing (so we can control the type name) @@ -136,45 +128,3 @@ def test_not_retryable_generic_exception(self): """Generic Exception without special name should NOT be retried.""" exc = Exception("generic error") assert _should_retry_exception(exc) is False - - -def _fake_response(): - choice = MagicMock() - choice.message.content = "ok" - choice.finish_reason = "stop" - resp = MagicMock() - resp.choices = [choice] - return resp - - -class TestRetryKwargForwarding: - """Regression tests for #233: the retry kwarg forwarded to LiteLLM must be - ``num_retries`` (LiteLLM's recognized internal control parameter), not - ``retries``. An unrecognized kwarg falls through as a provider - request-body field, which strict-mode proxies reject. - """ - - def test_llm_call_forwards_num_retries_not_retries(self): - with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() - ) as completion: - _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") - assert completion.call_args.kwargs["num_retries"] == 2 - assert "retries" not in completion.call_args.kwargs - - def test_llm_call_does_not_override_explicit_num_retries(self): - with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() - ) as completion: - _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step", num_retries=5) - assert completion.call_args.kwargs["num_retries"] == 5 - - def test_llm_call_async_forwards_num_retries_not_retries(self): - with patch( - "openkb.agent.compiler.litellm.acompletion", - new_callable=AsyncMock, - return_value=_fake_response(), - ) as acompletion: - asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step")) - assert acompletion.call_args.kwargs["num_retries"] == 2 - assert "retries" not in acompletion.call_args.kwargs From 1e90a7a0bff403126369a5e36c7744b361b0d2cc Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 16:07:40 +0200 Subject: [PATCH 08/16] revert: back out original fix/issue-229-retry-timeout merge (68ee872) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The branch was merged into integration twice: once early with the initial (buggy) 'retries=2' kwarg (68ee872, this revert), and again later with the 'num_retries' rename follow-up (4a4a3c1, already reverted). The prior revert of 4a4a3c1 alone left the original _should_retry_exception / kwargs.setdefault('retries', 2) code from 68ee872 still in place — completing the withdrawal here so integration carries no retry-classification code at all, matching the decision to withdraw #229/#230 pending root-cause investigation (see prior two revert commits). --- openkb/agent/compiler.py | 128 ++-------------------------------- tests/test_compiler_retry.py | 130 ----------------------------------- 2 files changed, 4 insertions(+), 254 deletions(-) delete mode 100644 tests/test_compiler_retry.py diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index 9ec19c32f..d0c9f878d 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -260,71 +260,6 @@ # --------------------------------------------------------------------------- -def _should_retry_exception(exc: Exception) -> bool: - """Determine whether an exception is retryable (transient error). - - Returns True for temporary API/network errors that may succeed on retry: - - Timeout (client-side or server-side) - - APIError 5xx (server errors) - - RateLimitError (429) - - ConnectionError / ServiceUnavailableError - - Returns False for permanent errors that won't be fixed by retry: - - TruncatedResponseError (model hit max_tokens) - - ValueError, TypeError (malformed input/output) - - AuthenticationError (credentials issue) - - BadRequestError (invalid parameters) - - Unknown error types (conservative approach) - """ - exc_type_name = type(exc).__name__ - - # ===== RETRYABLE (transient errors) ===== - - # Timeout (network/gateway timeout) - if "Timeout" in exc_type_name: - return True - - # Generic API errors (5xx range, but not 4xx) - if "APIError" in exc_type_name: - # Don't retry if it's a BadRequest/Invalid error (4xx) - if "Invalid" not in exc_type_name and "BadRequest" not in exc_type_name: - return True - - # Rate limiting (429) - if "RateLimitError" in exc_type_name or "Rate" in exc_type_name: - return True - - # Connection errors - if "ConnectionError" in exc_type_name: - return True - - # Service unavailable - if "ServiceUnavailable" in exc_type_name: - return True - - # ===== NOT RETRYABLE (permanent errors) ===== - - # Model hit max_tokens limit - if isinstance(exc, TruncatedResponseError): - return False - - # Content validation failures - if "ValueError" in exc_type_name or "TypeError" in exc_type_name: - return False - - # Authentication failures - if "Auth" in exc_type_name or "Permission" in exc_type_name: - return False - - # Bad parameters/requests - if "BadRequest" in exc_type_name or "Invalid" in exc_type_name: - return False - - # ===== UNKNOWN: Conservative approach ===== - # Don't retry errors we don't recognize - return False - - def _cached_text(text: str) -> list[dict]: """Wrap a text payload into a content-block list with an Anthropic ephemeral cache_control marker. @@ -471,11 +406,7 @@ def _llm_call( bundle=None, **kwargs, ) -> str: - """Single LLM call with animated progress, debug logging, and retry support. - - Transient errors (Timeout, 5xx, 429) are automatically retried by LiteLLM. - Permanent errors (4xx, truncation, validation) are raised immediately. - """ + """Single LLM call with animated progress and debug logging.""" messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -486,10 +417,6 @@ def _llm_call( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) - - # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) - kwargs.setdefault("retries", 2) - logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -498,27 +425,7 @@ def _llm_call( spinner.start() t0 = time.time() - try: - response = litellm.completion(model=model, messages=messages, **kwargs) - except Exception as exc: - # NEW: Better error logging with retry context - if _should_retry_exception(exc): - logger.warning( - "LLM [%s] failed with transient error (retries applied by LiteLLM): %s", - step_name, - exc, - exc_info=False, # Don't spam stack traces for known transient errors - ) - else: - logger.warning( - "LLM [%s] failed with permanent error (no retry): %s", - step_name, - exc, - exc_info=True, # Full trace for unexpected errors - ) - spinner.stop("[FAILED]") - raise - + response = litellm.completion(model=model, messages=messages, **kwargs) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) @@ -542,11 +449,7 @@ async def _llm_call_async( bundle=None, **kwargs, ) -> str: - """Async LLM call with timing output, debug logging, and retry support. - - Transient errors (Timeout, 5xx, 429) are automatically retried by LiteLLM. - Permanent errors (4xx, truncation, validation) are raised immediately. - """ + """Async LLM call with timing output and debug logging.""" messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -557,36 +460,13 @@ async def _llm_call_async( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) - - # NEW: Retry configuration for transient errors (fixed: 2 retries, base 2) - kwargs.setdefault("retries", 2) - logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) t0 = time.time() - try: - response = await litellm.acompletion(model=model, messages=messages, **kwargs) - except Exception as exc: - # NEW: Better error logging with retry context - if _should_retry_exception(exc): - logger.warning( - "LLM [%s] failed with transient error (retries applied by LiteLLM): %s", - step_name, - exc, - exc_info=False, - ) - else: - logger.warning( - "LLM [%s] failed with permanent error (no retry): %s", - step_name, - exc, - exc_info=True, - ) - raise - + response = await litellm.acompletion(model=model, messages=messages, **kwargs) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) diff --git a/tests/test_compiler_retry.py b/tests/test_compiler_retry.py deleted file mode 100644 index 7902201f7..000000000 --- a/tests/test_compiler_retry.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Tests for LLM retry logic in compiler.py.""" - -from openkb.agent.compiler import TruncatedResponseError, _should_retry_exception - - -# Custom exception classes for testing (so we can control the type name) -class TimeoutError(Exception): - """Simulates litellm.Timeout.""" - - pass - - -class APIError(Exception): - """Simulates litellm.APIError (5xx).""" - - pass - - -class InvalidAPIError(APIError): - """Simulates InvalidAPIError (4xx).""" - - pass - - -class BadRequestError(Exception): - """Simulates BadRequestError.""" - - pass - - -class RateLimitError(Exception): - """Simulates litellm.RateLimitError.""" - - pass - - -class AuthenticationError(Exception): - """Simulates AuthenticationError.""" - - pass - - -class PermissionError(Exception): - """Simulates PermissionError.""" - - pass - - -class ServiceUnavailableError(Exception): - """Simulates ServiceUnavailableError.""" - - pass - - -class TestShouldRetryException: - """Test the exception filtering logic for retry decisions.""" - - def test_retryable_timeout(self): - """Timeout should be retryable.""" - exc = TimeoutError("Gateway Timeout") - assert _should_retry_exception(exc) is True - - def test_retryable_api_error_5xx(self): - """5xx API errors should be retryable.""" - exc = APIError("503 Service Unavailable") - assert _should_retry_exception(exc) is True - - def test_not_retryable_invalid_api_error(self): - """InvalidAPIError (4xx) should NOT be retryable.""" - exc = InvalidAPIError("400 Bad Request") - assert _should_retry_exception(exc) is False - - def test_retryable_rate_limit(self): - """Rate limit errors should be retryable.""" - exc = RateLimitError("429 Too Many Requests") - assert _should_retry_exception(exc) is True - - def test_retryable_connection_error(self): - """Connection errors should be retryable.""" - exc = ConnectionError("Connection refused") - assert _should_retry_exception(exc) is True - - def test_retryable_service_unavailable(self): - """Service unavailable errors should be retryable.""" - exc = ServiceUnavailableError("Service down") - assert _should_retry_exception(exc) is True - - def test_not_retryable_truncation(self): - """Truncated output should NOT be retryable.""" - exc = TruncatedResponseError("hit length limit") - assert _should_retry_exception(exc) is False - - def test_not_retryable_value_error(self): - """ValueError should NOT be retryable.""" - exc = ValueError("empty content") - assert _should_retry_exception(exc) is False - - def test_not_retryable_type_error(self): - """TypeError should NOT be retryable.""" - exc = TypeError("malformed") - assert _should_retry_exception(exc) is False - - def test_not_retryable_auth_error(self): - """Authentication errors should NOT be retryable.""" - exc = AuthenticationError("invalid API key") - assert _should_retry_exception(exc) is False - - def test_not_retryable_permission_error(self): - """Permission errors should NOT be retryable.""" - exc = PermissionError("forbidden") - assert _should_retry_exception(exc) is False - - def test_not_retryable_bad_request(self): - """BadRequest errors should NOT be retryable.""" - exc = BadRequestError("invalid params") - assert _should_retry_exception(exc) is False - - def test_not_retryable_unknown(self): - """Unknown errors should NOT be retried (conservative).""" - - class WeirdCustomError(Exception): - pass - - exc = WeirdCustomError("something weird") - assert _should_retry_exception(exc) is False - - def test_not_retryable_generic_exception(self): - """Generic Exception without special name should NOT be retried.""" - exc = Exception("generic error") - assert _should_retry_exception(exc) is False From 941ed7b934c75be8bb02565861db0f95badb758c Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 16:25:44 +0200 Subject: [PATCH 09/16] feat(agent): add hybrid BM25 search_wiki tool to query/chat agent Adds a dependency-free BM25 full-text index (openkb/fulltext_index.py) over concepts/entities/summaries pages, exposed as a new search_wiki tool alongside index.md-driven navigation in build_query_agent. Additive hybrid retrieval: surfaces pages whose one-line index summary omits a buried detail, without replacing existing navigation. Resolves #233. --- README.md | 2 + openkb/agent/query.py | 29 +++++- openkb/agent/tools.py | 30 ++++++ openkb/fulltext_index.py | 177 +++++++++++++++++++++++++++++++++++ tests/test_agent_tools.py | 39 ++++++++ tests/test_fulltext_index.py | 103 ++++++++++++++++++++ tests/test_query.py | 3 +- 7 files changed, 378 insertions(+), 5 deletions(-) create mode 100644 openkb/fulltext_index.py create mode 100644 tests/test_fulltext_index.py diff --git a/README.md b/README.md index 988bebda0..0a5dbce82 100644 --- a/README.md +++ b/README.md @@ -207,6 +207,8 @@ A "generator" reads from the compiled wiki and produces something usable: an ans `openkb query "..."` answers a single question with a grounded, cited answer from your wiki. `openkb chat` is interactive, an ongoing multi-turn session over the same wiki (`--resume`, `--list`, `--delete` to manage sessions). → Walked through with real saved output in **[`examples/commands/`](examples/commands/)** (query) and **[`examples/chat/`](examples/chat/)** (chat). +Retrieval is hybrid: the agent primarily navigates via `index.md`'s one-line summaries, and additionally has a `search_wiki` tool — a dependency-free BM25 full-text search over `concepts/`, `entities/`, and `summaries/` — for surfacing pages whose index summary doesn't mention a specific buried detail. It's additive, not a replacement, so recall can only improve over index-only navigation. + Inside a chat, type `/` to access slash commands (Tab to complete).
diff --git a/openkb/agent/query.py b/openkb/agent/query.py index da1a939ef..e5cd79cd5 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -14,6 +14,9 @@ read_wiki_image, write_kb_file, ) +from openkb.agent.tools import ( + search_wiki as search_wiki_impl, +) from openkb.config import LlmCredentialBundle, resolve_model_settings from openkb.schema import get_agents_md @@ -33,18 +36,23 @@ 3. Read concept pages (concepts/) for cross-document synthesis. 4. For "who/what is X" questions about a specific named person, organization, place, or product, read the matching page in entities/ first. -5. When you need detailed source document content, each summary page has a +5. If index.md's one-line summaries don't surface a specific detail you + need (a niche term, an exact figure, a buried fact), use + search_wiki(query) — a keyword-level full-text search over + concepts/entities/summaries. This is a hybrid fallback: use it in + addition to, not instead of, index.md navigation. +6. When you need detailed source document content, each summary page has a `full_text` frontmatter field with the path to the original document content: - Short documents (doc_type: short): read_file with that path. - PageIndex documents (doc_type: pageindex): use get_page_content(doc_name, pages) with tight page ranges. The summary shows document tree structure with page ranges to help you target. Never fetch the whole document. -6. Source content may reference images. Short-doc .md pages link them +7. Source content may reference images. Short-doc .md pages link them note-relative (e.g. ![image](images/doc/file.png), resolved from wiki/sources/); long-doc JSON page metadata lists them wiki-root-relative (e.g. sources/images/doc/file.png). Pass either form as seen to the get_image tool — it accepts both. -7. Synthesize a clear, concise, well-cited answer grounded in wiki content. +8. Synthesize a clear, concise, well-cited answer grounded in wiki content. Answer based only on wiki content. Be concise. Before each tool call, output one short sentence explaining the reason. @@ -83,6 +91,19 @@ def get_page_content(doc_name: str, pages: str) -> str: """ return get_wiki_page_content(doc_name, pages, wiki_root) + @function_tool + def search_wiki(query: str) -> str: + """Full-text (BM25) keyword search over concepts/entities/summaries. + + Hybrid fallback for when index.md's one-line summaries don't surface + a specific buried detail (a niche term, an exact figure, a fact). + Use in addition to, not instead of, index.md navigation. + + Args: + query: Free-text search query (keywords or a natural-language question). + """ + return search_wiki_impl(query, wiki_root) + @function_tool def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: """View an image from the wiki. @@ -117,7 +138,7 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: return Agent( name="wiki-query", instructions=instructions, - tools=[read_file, get_page_content, get_image], + tools=[read_file, get_page_content, search_wiki, get_image], model=f"litellm/{model}", model_settings=ModelSettings(**model_settings), ) diff --git a/openkb/agent/tools.py b/openkb/agent/tools.py index eedd388de..a4fa3ad91 100644 --- a/openkb/agent/tools.py +++ b/openkb/agent/tools.py @@ -135,6 +135,36 @@ def get_wiki_page_content(doc_name: str, pages: str, wiki_root: str) -> str: return "\n\n".join(parts) + "\n\n" +def search_wiki(query: str, wiki_root: str, top_k: int = 5) -> str: + """Full-text (BM25) search over concepts/entities/summaries wiki pages. + + Hybrid retrieval helper: complements index.md-driven navigation by + surfacing pages whose one-line index summary doesn't mention a specific + buried detail the query is looking for (a niche term, a figure, an exact + fact). Additive — use alongside, not instead of, index.md navigation. + + Args: + query: Free-text search query (keywords or a natural-language question). + wiki_root: Absolute path to the wiki root directory. + top_k: Maximum number of ranked results to return. + + Returns: + A formatted, ranked list of page hits (wikilink, title, snippet), or + a message indicating no matches were found. + """ + from openkb.fulltext_index import WikiFullTextIndex + + hits = WikiFullTextIndex(wiki_root).search(query, top_k=top_k) + if not hits: + return "No matching pages found." + + lines = [] + for i, hit in enumerate(hits, start=1): + wikilink = hit.path[:-3] if hit.path.endswith(".md") else hit.path + lines.append(f"{i}. [[{wikilink}]] — {hit.title} (score: {hit.score})\n {hit.snippet}") + return "\n".join(lines) + + _MIME_TYPES = { ".png": "image/png", ".jpg": "image/jpeg", diff --git a/openkb/fulltext_index.py b/openkb/fulltext_index.py new file mode 100644 index 000000000..5c7852df9 --- /dev/null +++ b/openkb/fulltext_index.py @@ -0,0 +1,177 @@ +"""Dependency-free BM25 full-text index over compiled wiki pages. + +Hybrid retrieval: the query/chat agent's primary search strategy is +``index.md`` navigation (one-line summaries pointing at pages to read). That +strategy loses recall for details buried deep in a page body that the +one-liner doesn't mention. This module adds an additive, keyword-level +fallback — a BM25 index over the same compiled pages — exposed to the agent +as the ``search_wiki`` tool (see ``openkb.agent.tools.search_wiki``). It is a +union with index-driven navigation, not a replacement, so recall can only +improve relative to index-only navigation, never regress. + +No new dependency: OpenKB pins dependencies exactly and vets each one +deliberately (see ``pyproject.toml``), and BM25 over a few hundred wiki pages +is cheap enough in pure Python that a search-library dependency (e.g. Whoosh) +isn't warranted. +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from pathlib import Path + +from openkb.schema import PAGE_CONTENT_DIRS + +_TOKEN_RE = re.compile(r"[a-z0-9]+") + +# Standard BM25 hyperparameters (Robertson/Sparck-Jones defaults). +_K1 = 1.5 +_B = 0.75 + +_SNIPPET_RADIUS = 80 # characters of context on each side of the first match + + +def _tokenize(text: str) -> list[str]: + """Lowercase, alphanumeric-only tokenization (no stemming).""" + return _TOKEN_RE.findall(text.lower()) + + +def _extract_title(text: str) -> str | None: + """Return the first ``# heading`` line's text, or ``None``.""" + for line in text.splitlines(): + stripped = line.strip() + if stripped.startswith("# "): + return stripped[2:].strip() + return None + + +def _make_snippet(text: str, query_terms: list[str]) -> str: + """Return a short excerpt around the first query-term match in *text*.""" + lowered = text.lower() + match_pos = -1 + for term in query_terms: + pos = lowered.find(term) + if pos != -1 and (match_pos == -1 or pos < match_pos): + match_pos = pos + if match_pos == -1: + collapsed = " ".join(text.split()) + truncated = collapsed[: _SNIPPET_RADIUS * 2] + suffix = "…" if len(collapsed) > _SNIPPET_RADIUS * 2 else "" + return truncated + suffix + + start = max(0, match_pos - _SNIPPET_RADIUS) + end = min(len(text), match_pos + _SNIPPET_RADIUS) + collapsed = " ".join(text[start:end].split()) + prefix = "…" if start > 0 else "" + suffix = "…" if end < len(text) else "" + return f"{prefix}{collapsed}{suffix}" + + +@dataclass(frozen=True) +class SearchHit: + """A single BM25 search result over a wiki page.""" + + path: str # wiki-root-relative, e.g. "concepts/attention.md" + title: str + score: float + snippet: str + + +@dataclass(frozen=True) +class _IndexedPage: + path: str + title: str + text: str + tokens: list[str] + + +class WikiFullTextIndex: + """In-memory BM25 index over :data:`PAGE_CONTENT_DIRS` wiki pages. + + Rebuilt fresh on construction — cheap enough at the wiki sizes this + pattern targets (hundreds of pages); no on-disk cache or incremental + update is needed. + """ + + def __init__(self, wiki_root: str | Path) -> None: + self._wiki_root = Path(wiki_root).resolve() + self._pages: list[_IndexedPage] = [] + self._df: dict[str, int] = {} + self._avgdl = 0.0 + self._build() + + def _build(self) -> None: + for subdir in PAGE_CONTENT_DIRS: + target = self._wiki_root / subdir + if not target.is_dir(): + continue + for md_file in sorted(target.glob("*.md")): + text = md_file.read_text(encoding="utf-8") + tokens = _tokenize(text) + if not tokens: + continue + title = _extract_title(text) or md_file.stem + path = f"{subdir}/{md_file.name}" + self._pages.append(_IndexedPage(path=path, title=title, text=text, tokens=tokens)) + + if not self._pages: + return + + self._avgdl = sum(len(page.tokens) for page in self._pages) / len(self._pages) + for page in self._pages: + for term in set(page.tokens): + self._df[term] = self._df.get(term, 0) + 1 + + def _idf(self, term: str) -> float: + n = len(self._pages) + df = self._df.get(term, 0) + # +1 smoothing keeps idf non-negative even for very common terms. + return math.log((n - df + 0.5) / (df + 0.5) + 1) + + def _score(self, query_terms: list[str], page: _IndexedPage) -> float: + dl = len(page.tokens) + tf: dict[str, int] = {} + for term in page.tokens: + tf[term] = tf.get(term, 0) + 1 + + score = 0.0 + for term in query_terms: + f = tf.get(term, 0) + if f == 0: + continue + idf = self._idf(term) + numerator = f * (_K1 + 1) + denominator = f + _K1 * (1 - _B + _B * dl / self._avgdl) + score += idf * (numerator / denominator) + return score + + def search(self, query: str, top_k: int = 5) -> list[SearchHit]: + """Return the ``top_k`` highest-scoring pages for *query* (BM25). + + Args: + query: Free-text search query (keywords or a question). + top_k: Maximum number of results to return. + + Returns: + Ranked hits, highest score first. Empty if the query has no + tokens or the index has no pages. + """ + query_terms = _tokenize(query) + if not query_terms or not self._pages: + return [] + + scored = [(self._score(query_terms, page), page) for page in self._pages] + scored = [(score, page) for score, page in scored if score > 0] + scored.sort(key=lambda item: item[0], reverse=True) + + return [ + SearchHit( + path=page.path, + title=page.title, + score=round(score, 3), + snippet=_make_snippet(page.text, query_terms), + ) + for score, page in scored[:top_k] + ] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 283a5a8b0..9c0463629 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -9,6 +9,7 @@ parse_pages, read_wiki_file, read_wiki_image, + search_wiki, write_wiki_file, ) @@ -320,3 +321,41 @@ def test_artifact_event_none_for_non_output_zone(): def test_artifact_event_none_for_bad_json(): assert artifact_event_from_write("write_file", "not json", "Written: output/x.html") is None + + +# --------------------------------------------------------------------------- +# search_wiki +# --------------------------------------------------------------------------- + + +class TestSearchWiki: + def test_finds_matching_page(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "cnn.md").write_text( + "# Convolutional Neural Networks\n\nDropout regularization prevents overfitting." + ) + + result = search_wiki("dropout regularization", wiki_root) + + assert "[[concepts/cnn]]" in result + assert "Convolutional Neural Networks" in result + + def test_no_matches_returns_message(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "cnn.md").write_text("# CNN\n\nSomething else entirely.") + + result = search_wiki("nonexistent_keyword_xyz", wiki_root) + + assert result == "No matching pages found." + + def test_respects_top_k(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "entities").mkdir() + for i in range(5): + (tmp_path / "entities" / f"e{i}.md").write_text(f"# Entity {i}\n\nkeyword {i}.") + + result = search_wiki("keyword", wiki_root, top_k=2) + + assert result.count("[[entities/") == 2 diff --git a/tests/test_fulltext_index.py b/tests/test_fulltext_index.py new file mode 100644 index 000000000..750b7498b --- /dev/null +++ b/tests/test_fulltext_index.py @@ -0,0 +1,103 @@ +"""Tests for openkb.fulltext_index (BM25 hybrid search).""" + +from __future__ import annotations + +from openkb.fulltext_index import WikiFullTextIndex + + +def _write(tmp_path, subdir, name, text): + directory = tmp_path / subdir + directory.mkdir(parents=True, exist_ok=True) + (directory / name).write_text(text, encoding="utf-8") + + +class TestWikiFullTextIndex: + def test_empty_wiki_returns_no_hits(self, tmp_path): + index = WikiFullTextIndex(str(tmp_path)) + assert index.search("anything") == [] + + def test_finds_page_by_keyword_in_body(self, tmp_path): + _write( + tmp_path, + "concepts", + "cnn.md", + "# Convolutional Neural Networks\n\nAlexNet popularized ReLU activations " + "and dropout regularization for large-scale image classification.", + ) + _write( + tmp_path, + "concepts", + "unrelated.md", + "# Gardening\n\nTomatoes need plenty of sunlight and water.", + ) + + hits = WikiFullTextIndex(str(tmp_path)).search("dropout regularization") + + assert len(hits) == 1 + assert hits[0].path == "concepts/cnn.md" + assert hits[0].title == "Convolutional Neural Networks" + assert hits[0].score > 0 + + def test_ranks_more_relevant_page_higher(self, tmp_path): + _write( + tmp_path, + "concepts", + "on-topic.md", + "# Topic\n\nAlexNet AlexNet AlexNet training data criticism bias bias.", + ) + _write( + tmp_path, + "concepts", + "off-topic.md", + "# Other\n\nA single passing mention of AlexNet in an unrelated paragraph " + "about something else entirely, padded with filler words to change length.", + ) + + hits = WikiFullTextIndex(str(tmp_path)).search("AlexNet bias") + + assert [hit.path for hit in hits[:1]] == ["concepts/on-topic.md"] + + def test_respects_top_k(self, tmp_path): + for i in range(10): + _write(tmp_path, "entities", f"e{i}.md", f"# Entity {i}\n\nkeyword appears here {i}.") + + hits = WikiFullTextIndex(str(tmp_path)).search("keyword", top_k=3) + + assert len(hits) == 3 + + def test_only_indexes_page_content_dirs(self, tmp_path): + _write(tmp_path, "sources", "raw.md", "# Raw\n\nkeyword raw source content.") + _write(tmp_path, "concepts", "c.md", "# Concept\n\nkeyword concept content.") + + hits = WikiFullTextIndex(str(tmp_path)).search("keyword") + + assert [hit.path for hit in hits] == ["concepts/c.md"] + + def test_falls_back_to_filename_when_no_heading(self, tmp_path): + _write(tmp_path, "summaries", "no-heading.md", "keyword content without a heading line.") + + hits = WikiFullTextIndex(str(tmp_path)).search("keyword") + + assert hits[0].title == "no-heading" + + def test_no_query_tokens_returns_no_hits(self, tmp_path): + _write(tmp_path, "concepts", "c.md", "# Concept\n\nkeyword concept content.") + + hits = WikiFullTextIndex(str(tmp_path)).search(" ") + + assert hits == [] + + def test_snippet_contains_context_around_match(self, tmp_path): + _write( + tmp_path, + "concepts", + "c.md", + "# Concept\n\n" + + ("padding " * 40) + + "the exact fee is five hundred dollars" + + (" more" * 40), + ) + + hits = WikiFullTextIndex(str(tmp_path)).search("fee") + + assert "fee" in hits[0].snippet.lower() diff --git a/tests/test_query.py b/tests/test_query.py index ecaceabd9..a720ccce3 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -19,13 +19,14 @@ def test_agent_name(self, tmp_path): def test_agent_has_three_tools(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") - assert len(agent.tools) == 3 + assert len(agent.tools) == 4 def test_agent_tool_names(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") names = {t.name for t in agent.tools} assert "read_file" in names assert "get_page_content" in names + assert "search_wiki" in names assert "get_image" in names def test_instructions_mention_get_page_content(self, tmp_path): From c7ea8ab546bd412df1d17f77f3312910ff4313e5 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 17:23:42 +0200 Subject: [PATCH 10/16] fix(agent): stream LLM completions to avoid gateway idle-timeout Corporate LLM gateways (e.g. AI.proxy on AWS) enforce an idle timeout on buffered (non-streaming) requests, so a long-running compile step can hit a Gateway Timeout even though the provider would have eventually finished. Switch _llm_call() and _llm_call_async() in openkb/agent/compiler.py to litellm.completion()/acompletion() with stream=True: streaming keeps bytes flowing over the connection, so idle-timeout gateways never see a silent connection. Chunks are merged back into the existing response shape via a new _merge_stream_chunks() helper, using LiteLLM's own litellm.stream_chunk_builder() for genuine multi-chunk streams. An exception raised mid-stream propagates as a complete failure (list() never returns a partial buffer), matching prior all-or-nothing behavior. Adapts the compiler test mocks (_mock_completion/_mock_acompletion and a handful of inline mocks) to return a single-chunk fake stream, plus the litellm.completion/acompletion mocks in test_llm_timeout.py. Resolves #235. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- openkb/agent/compiler.py | 49 ++++++++++++++++-- tests/test_compiler.py | 106 ++++++++++++-------------------------- tests/test_llm_timeout.py | 19 ++++--- 3 files changed, 91 insertions(+), 83 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878d..2f9df71e1 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -397,6 +397,23 @@ class TruncatedResponseError(Exception): treat truncation as a failure (so a partial page is skipped, not written).""" +def _merge_stream_chunks(chunks: list, messages: list[dict]): + """Merge streamed LLM chunks back into a single, non-streaming response. + + Genuine LiteLLM stream chunks only ever carry a ``.delta`` (never a + ``.message``), so a real multi-chunk stream is merged via LiteLLM's own + :func:`litellm.stream_chunk_builder`. A single chunk that already looks + like a complete, non-streaming ``ModelResponse`` (exposing ``.message``) + is used as-is — there's nothing left to merge, and it lets test doubles + fake a one-shot response without simulating LiteLLM's internal delta + format. + """ + choices = getattr(chunks[0], "choices", None) or [] + if len(chunks) == 1 and choices and hasattr(choices[0], "message"): + return chunks[0] + return litellm.stream_chunk_builder(chunks, messages=messages) + + def _llm_call( model: str, messages: list[dict], @@ -406,7 +423,15 @@ def _llm_call( bundle=None, **kwargs, ) -> str: - """Single LLM call with animated progress and debug logging.""" + """Single LLM call with animated progress and debug logging. + + Uses ``stream=True``: some corporate LLM gateways enforce an idle + timeout on buffered (non-streaming) requests, which a long-running + completion can hit before the response is ever sent. Streaming keeps + bytes flowing over the connection so that timeout never fires; the + chunks are merged back into a single response via + :func:`_merge_stream_chunks` so callers see the same shape as before. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -417,6 +442,7 @@ def _llm_call( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + kwargs.setdefault("stream_options", {"include_usage": True}) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) @@ -425,7 +451,11 @@ def _llm_call( spinner.start() t0 = time.time() - response = litellm.completion(model=model, messages=messages, **kwargs) + stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs) + chunks = list(stream) + if not chunks: + raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") + response = _merge_stream_chunks(chunks, messages) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) @@ -449,7 +479,10 @@ async def _llm_call_async( bundle=None, **kwargs, ) -> str: - """Async LLM call with timing output and debug logging.""" + """Async LLM call with timing output and debug logging. + + See ``_llm_call`` for why ``stream=True`` is used. + """ messages = _prepare_messages(model, messages) extra_headers = bundle.extra_headers if bundle is not None else get_extra_headers() if extra_headers: @@ -460,13 +493,21 @@ async def _llm_call_async( if bundle is not None: kwargs.setdefault("api_key", bundle.api_key) kwargs.setdefault("base_url", bundle.base_url) + kwargs.setdefault("stream_options", {"include_usage": True}) logger.debug("LLM request [%s]:\n%s", step_name, _fmt_messages(messages)) if kwargs: logger.debug("LLM kwargs [%s]: %s", step_name, kwargs) t0 = time.time() - response = await litellm.acompletion(model=model, messages=messages, **kwargs) + stream = await litellm.acompletion(model=model, messages=messages, stream=True, **kwargs) + if hasattr(stream, "__aiter__"): + chunks = [chunk async for chunk in stream] + else: + chunks = list(stream) + if not chunks: + raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") + response = _merge_stream_chunks(chunks, messages) content = response.choices[0].message.content or "" truncated = _warn_if_truncated(response, step_name, kwargs.get("max_tokens")) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 95a57cc4c..e17ad47d7 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1112,36 +1112,45 @@ def test_frontmatter_without_sources_line_gets_one_inserted(self, tmp_path): assert "[[summaries/new-doc]]" in text +def _mock_response(content, finish_reason: str = "stop") -> MagicMock: + """Build a fake, already-complete LLM response (single-chunk stream). + + ``_llm_call``/``_llm_call_async`` now call ``litellm.completion``/ + ``acompletion`` with ``stream=True`` and merge the resulting chunks back + into one response (see ``_merge_stream_chunks``). Exposing ``.message`` + (rather than the ``.delta`` a genuine stream chunk carries) tells + ``_merge_stream_chunks`` this single chunk *is* the final response, so it + is used as-is without needing to fake LiteLLM's internal delta format. + """ + mock_resp = MagicMock() + mock_resp.choices = [MagicMock()] + mock_resp.choices[0].message.content = content + mock_resp.choices[0].finish_reason = finish_reason + mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) + mock_resp.usage.prompt_tokens_details = None + return mock_resp + + def _mock_completion(responses: list[str]): - """Create a mock for litellm.completion that returns responses in order.""" + """Create a mock for litellm.completion returning a single-chunk stream.""" call_count = {"n": 0} def side_effect(*args, **kwargs): idx = min(call_count["n"], len(responses) - 1) call_count["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(responses[idx])] return side_effect def _mock_acompletion(responses: list[str]): - """Create an async mock for litellm.acompletion.""" + """Create an async mock for litellm.acompletion returning a single-chunk stream.""" call_count = {"n": 0} async def side_effect(*args, **kwargs): idx = min(call_count["n"], len(responses) - 1) call_count["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(responses[idx])] return side_effect @@ -1342,15 +1351,7 @@ def sync_side_effect(*args, **kwargs): sync_call_count["n"] += 1 if idx == 2: # the summary-rewrite call raise RuntimeError("simulated API failure") - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = [ - summary_response, - plan_response, - ][idx] - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response([summary_response, plan_response][idx])] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1507,21 +1508,11 @@ async def test_short_doc_marks_doc_and_summary(self, tmp_path): def sync_side_effect(*args, **kwargs): captured_sync_calls.append(kwargs["messages"]) idx = min(len(captured_sync_calls) - 1, len(sync_responses) - 1) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = sync_responses[idx] - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(sync_responses[idx])] async def async_side_effect(*args, **kwargs): captured_async_calls.append(kwargs["messages"]) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = concept_response - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(concept_response)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1586,15 +1577,9 @@ async def test_long_doc_marks_doc_message(self, tmp_path): def sync_side_effect(*args, **kwargs): captured.append(kwargs["messages"]) - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] # First call: overview (plain text); second: plan (JSON). - mock_resp.choices[0].message.content = ( - "Overview text" if len(captured) == 1 else plan_response - ) - mock_resp.usage = MagicMock(prompt_tokens=1, completion_tokens=1) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + content = "Overview text" if len(captured) == 1 else plan_response + return [_mock_response(content)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=sync_side_effect) @@ -1726,16 +1711,9 @@ async def test_create_and_update_flow(self, tmp_path): async def ordered_acompletion(*args, **kwargs): idx = call_order["n"] call_order["n"] += 1 - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] # create tasks come first, then update tasks - if idx == 0: - mock_resp.choices[0].message.content = create_page_response - else: - mock_resp.choices[0].message.content = update_page_response - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + content = create_page_response if idx == 0 else update_page_response + return [_mock_response(content)] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1823,13 +1801,7 @@ async def test_truncated_update_preserves_existing_page(self, tmp_path): ) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1859,13 +1831,7 @@ async def test_truncated_create_skips_partial_page(self, tmp_path): truncated_page = json.dumps({"brief": "x", "content": "# Ghost\n\nPartial"}) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) @@ -1928,13 +1894,7 @@ async def test_truncated_entity_update_preserves_existing_page(self, tmp_path): ) async def truncated_acompletion(*args, **kwargs): - mock_resp = MagicMock() - mock_resp.choices = [MagicMock()] - mock_resp.choices[0].message.content = truncated_page - mock_resp.choices[0].finish_reason = "length" - mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) - mock_resp.usage.prompt_tokens_details = None - return mock_resp + return [_mock_response(truncated_page, finish_reason="length")] with patch("openkb.agent.compiler.litellm") as mock_litellm: mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) diff --git a/tests/test_llm_timeout.py b/tests/test_llm_timeout.py index ca7d80e68..db119df3c 100644 --- a/tests/test_llm_timeout.py +++ b/tests/test_llm_timeout.py @@ -17,6 +17,13 @@ def _fake_response(): + """A fake, already-complete LLM response (single-chunk stream). + + See ``openkb.agent.compiler._merge_stream_chunks``: a chunk exposing + ``.message`` (as this one does) is treated as already-complete and used + as-is, so callers of ``litellm.completion``/``acompletion`` with + ``stream=True`` can be mocked to just return a one-item list. + """ choice = MagicMock() choice.message.content = "ok" choice.finish_reason = "stop" @@ -28,7 +35,7 @@ def _fake_response(): def test_llm_call_forwards_configured_timeout(): set_timeout(1200.0) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()] ) as completion: _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") assert completion.call_args.kwargs["timeout"] == 1200.0 @@ -37,7 +44,7 @@ def test_llm_call_forwards_configured_timeout(): def test_llm_call_omits_timeout_when_unset(): set_timeout(None) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()] ) as completion: _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step") assert "timeout" not in completion.call_args.kwargs @@ -47,7 +54,7 @@ def test_llm_call_does_not_override_explicit_timeout(): # An explicit per-call timeout kwarg wins over the configured default. set_timeout(1200.0) with patch( - "openkb.agent.compiler.litellm.completion", return_value=_fake_response() + "openkb.agent.compiler.litellm.completion", return_value=[_fake_response()] ) as completion: _llm_call("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30) assert completion.call_args.kwargs["timeout"] == 30 @@ -58,7 +65,7 @@ def test_llm_call_async_forwards_configured_timeout(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + return_value=[_fake_response()], ) as acompletion: asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step")) assert acompletion.call_args.kwargs["timeout"] == 900.0 @@ -69,7 +76,7 @@ def test_llm_call_async_omits_timeout_when_unset(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + return_value=[_fake_response()], ) as acompletion: asyncio.run(_llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step")) assert "timeout" not in acompletion.call_args.kwargs @@ -80,7 +87,7 @@ def test_llm_call_async_does_not_override_explicit_timeout(): with patch( "openkb.agent.compiler.litellm.acompletion", new_callable=AsyncMock, - return_value=_fake_response(), + return_value=[_fake_response()], ) as acompletion: asyncio.run( _llm_call_async("gpt-4o", [{"role": "user", "content": "hi"}], "step", timeout=30) From 76905b8c90c4c006ff356e4b1e7a74b86a826dfc Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Sat, 29 Aug 2026 16:07:41 +0200 Subject: [PATCH 11/16] fix(agent): add debug timing for streamed completions Add per-chunk debug logging around streamed LiteLLM completion consumption in `openkb.agent.compiler`. This diagnostic instrumentation is enabled via the existing `openkb -v` flag and helps narrow the still-observed ~60s production cutoff to either a no-first-byte case (for example slow time-to-first-token) or a proxy/gateway path that buffers or drops streamed bytes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- openkb/agent/compiler.py | 86 ++++++++++++++++++++- tests/test_compiler.py | 162 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 245 insertions(+), 3 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index 2f9df71e1..311ce4f16 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -414,6 +414,86 @@ def _merge_stream_chunks(chunks: list, messages: list[dict]): return litellm.stream_chunk_builder(chunks, messages=messages) +def _log_chunk_timing( + step_name: str, chunk_number: int, t0: float, last_t: float, now: float +) -> None: + """Debug-log arrival timing for one streamed chunk. + + Run with ``openkb -v`` to see, per LLM call, how long the first chunk took + (time-to-first-token) and the gap to each subsequent chunk. If chunks do + arrive before a timeout, that shows the proxy/gateway is forwarding the + stream; if no chunk is ever logged before a timeout, the instrumentation + narrows the problem to a no-first-byte case (e.g. slow TTFT or an + intermediary buffering the response). + """ + logger.debug( + "LLM stream chunk [%s] #%d after %.2fs total (+%.2fs since previous)", + step_name, + chunk_number, + now - t0, + now - last_t, + ) + + +def _consume_stream(stream, step_name: str, t0: float) -> list: + """Collect a sync LiteLLM stream into a list, logging per-chunk timing. + + A mid-stream exception (e.g. the gateway idle-timeout firing) propagates + after logging how many chunks arrived and when, so callers still see a + complete failure — no partial buffer is ever returned. + """ + if not logger.isEnabledFor(logging.DEBUG): + return list(stream) + + chunks: list = [] + last_t = t0 + try: + for chunk in stream: + now = time.time() + _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now) + chunks.append(chunk) + last_t = now + except Exception: + logger.debug( + "LLM stream [%s] failed after %.2fs with %d chunk(s) received", + step_name, + time.time() - t0, + len(chunks), + exc_info=True, + ) + raise + return chunks + + +async def _consume_stream_async(stream, step_name: str, t0: float) -> list: + """Collect an async LiteLLM stream into a list, logging per-chunk timing. + + Mirrors :func:`_consume_stream`, including the no-partial-buffer invariant + on failure. + """ + if not logger.isEnabledFor(logging.DEBUG): + return [chunk async for chunk in stream] + + chunks: list = [] + last_t = t0 + try: + async for chunk in stream: + now = time.time() + _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now) + chunks.append(chunk) + last_t = now + except Exception: + logger.debug( + "LLM stream [%s] failed after %.2fs with %d chunk(s) received", + step_name, + time.time() - t0, + len(chunks), + exc_info=True, + ) + raise + return chunks + + def _llm_call( model: str, messages: list[dict], @@ -452,7 +532,7 @@ def _llm_call( t0 = time.time() stream = litellm.completion(model=model, messages=messages, stream=True, **kwargs) - chunks = list(stream) + chunks = _consume_stream(stream, step_name, t0) if not chunks: raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") response = _merge_stream_chunks(chunks, messages) @@ -502,9 +582,9 @@ async def _llm_call_async( stream = await litellm.acompletion(model=model, messages=messages, stream=True, **kwargs) if hasattr(stream, "__aiter__"): - chunks = [chunk async for chunk in stream] + chunks = await _consume_stream_async(stream, step_name, t0) else: - chunks = list(stream) + chunks = _consume_stream(stream, step_name, t0) if not chunks: raise RuntimeError(f"LLM [{step_name}] stream produced no chunks") response = _merge_stream_chunks(chunks, messages) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index e17ad47d7..a5119ff2a 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1155,6 +1156,53 @@ async def side_effect(*args, **kwargs): return side_effect +class _NoOpSpinner: + """Test double that disables spinner side effects.""" + + def __init__(self, *_args, **_kwargs): + pass + + def start(self) -> None: + pass + + def stop(self, _suffix: str = "") -> None: + pass + + +class _AsyncStream: + """Simple async iterator for exercising streamed LiteLLM responses in tests.""" + + def __init__( + self, + chunks: list[object], + *, + error: Exception | None = None, + raise_after: int | None = None, + ) -> None: + self._chunks = chunks + self._error = error + self._raise_after = raise_after + self._index = 0 + + def __aiter__(self) -> _AsyncStream: + return self + + async def __anext__(self) -> object: + if self._raise_after is not None and self._index == self._raise_after: + raise self._error or RuntimeError("stream exploded") + if self._index >= len(self._chunks): + raise StopAsyncIteration + chunk = self._chunks[self._index] + self._index += 1 + return chunk + + +def _stream_then_raise(chunks: list[object], error: Exception): + """Yield all chunks, then raise ``error`` on the next iteration.""" + yield from chunks + raise error + + class TestCompileShortDoc: @pytest.mark.asyncio async def test_full_pipeline(self, tmp_path): @@ -2661,6 +2709,120 @@ async def test_llm_call_async_injects_extra_headers(self): assert kwargs["extra_headers"] == {"Copilot-Integration-Id": "vscode-chat"} +class TestLLMStreamTimingDebugLogging: + """Per-chunk debug timing should be visible when verbose logging is enabled.""" + + def test_llm_call_logs_each_stream_chunk_at_debug(self, caplog): + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock( + return_value=iter(["chunk-1", "chunk-2", "chunk-3"]) + ) + mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok")) + + out = _llm_call("m", [{"role": "user", "content": "hi"}], "sync-step") + + assert out == "ok" + chunk_logs = [ + record.getMessage() + for record in caplog.records + if "LLM stream chunk [sync-step]" in record.getMessage() + ] + assert len(chunk_logs) == 3 + assert "#1" in chunk_logs[0] + assert "#3" in chunk_logs[-1] + + @pytest.mark.asyncio + async def test_llm_call_async_logs_each_stream_chunk_at_debug(self, caplog): + from openkb.agent.compiler import _llm_call_async + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.acompletion = AsyncMock( + return_value=_AsyncStream(["chunk-1", "chunk-2", "chunk-3"]) + ) + mock_litellm.stream_chunk_builder = MagicMock(return_value=_mock_response("ok")) + + out = await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-step") + + assert out == "ok" + chunk_logs = [ + record.getMessage() + for record in caplog.records + if "LLM stream chunk [async-step]" in record.getMessage() + ] + assert len(chunk_logs) == 3 + assert "#1" in chunk_logs[0] + assert "#3" in chunk_logs[-1] + + def test_llm_call_logs_stream_failure_before_reraising(self, caplog): + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("stream exploded") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock( + return_value=_stream_then_raise(["chunk-1", "chunk-2"], error) + ) + + with pytest.raises(RuntimeError, match="stream exploded"): + _llm_call("m", [{"role": "user", "content": "hi"}], "sync-fail-step") + + chunk_logs = [ + record.getMessage() + for record in caplog.records + if "LLM stream chunk [sync-fail-step]" in record.getMessage() + ] + failure_logs = [ + record.getMessage() + for record in caplog.records + if "LLM stream [sync-fail-step] failed after" in record.getMessage() + ] + assert len(chunk_logs) == 2 + assert len(failure_logs) == 1 + assert "2 chunk(s) received" in failure_logs[0] + + @pytest.mark.asyncio + async def test_llm_call_async_logs_stream_failure_before_reraising(self, caplog): + from openkb.agent.compiler import _llm_call_async + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("stream exploded") + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.acompletion = AsyncMock( + return_value=_AsyncStream( + ["chunk-1", "chunk-2"], + error=error, + raise_after=2, + ) + ) + + with pytest.raises(RuntimeError, match="stream exploded"): + await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-fail-step") + + chunk_logs = [ + record.getMessage() + for record in caplog.records + if "LLM stream chunk [async-fail-step]" in record.getMessage() + ] + failure_logs = [ + record.getMessage() + for record in caplog.records + if "LLM stream [async-fail-step] failed after" in record.getMessage() + ] + assert len(chunk_logs) == 2 + assert len(failure_logs) == 1 + assert "2 chunk(s) received" in failure_logs[0] + + class TestCacheControlStripping: """cache_control markers must only reach providers that honour them. From df469623e3d59b1199cf69b2b28ca3eccf9669b2 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Sun, 30 Aug 2026 11:07:54 +0200 Subject: [PATCH 12/16] feat(cli): redirect per-file debug logging to logs/.log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Console output during a batch openkb add used to interleave every DEBUG record (full LLM request/response dumps, tracebacks) with the normal status lines once -v was on, making a large batch unreadable. DEBUG detail is now written per processed file to logs/..log (created next to raw/, sources/, wiki/) instead, while the console keeps showing only the usual Adding/[OK]/[ERROR] status lines and any WARNING+ message. - _configure_console_logging(): pins the console handler's own level to WARNING, decoupled from whatever level a logger further up allows through, so DEBUG records can be created without leaking to stdout/ stderr. - _per_file_debug_log(): a context manager that, when enabled, raises the openkb logger to DEBUG and attaches a fresh FileHandler for the duration of one add_single_file() call, then detaches and restores the prior level — a no-op when debug logging isn't active. - New debug: true KB config.yaml key (DEFAULT_CONFIG) lets this be enabled per-KB without passing -v on every command; deliberately left out of GLOBAL_SCALAR_KEYS (no global.yaml default, no REST API/ workbench exposure) since it's a CLI-only diagnostic switch. Not yet an upstream issue/PR — local branch for now. --- config.yaml.example | 8 ++ examples/configuration/README.md | 25 ++++++ openkb/cli.py | 97 +++++++++++++++++++++-- openkb/config.py | 8 ++ tests/test_add_command.py | 129 +++++++++++++++++++++++++++++++ 5 files changed, 262 insertions(+), 5 deletions(-) diff --git a/config.yaml.example b/config.yaml.example index 47d2fda2b..8e691af5e 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -36,3 +36,11 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex # extra_headers: # extra HTTP headers some providers need (e.g. GitHub Copilot) # Editor-Version: vscode/1.95.0 # Copilot-Integration-Id: vscode-chat + +# Optional: enable detailed per-file debug logging without passing `-v` on +# every command. Each processed file gets its own log at +# logs/..log (created next to raw/, sources/, wiki/) with the +# full LLM request/response dump and tracebacks; the console keeps showing +# only the normal status lines and WARNING+ messages. Handy for unattended +# batch `add` runs where you only want to dig into the log after a failure. +# debug: true diff --git a/examples/configuration/README.md b/examples/configuration/README.md index 25e5b3a85..d9d889254 100644 --- a/examples/configuration/README.md +++ b/examples/configuration/README.md @@ -104,6 +104,14 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex # extra_headers: # extra HTTP headers some providers need (e.g. GitHub Copilot) # Editor-Version: vscode/1.95.0 # Copilot-Integration-Id: vscode-chat + +# Optional: enable detailed per-file debug logging without passing `-v` on +# every command. Each processed file gets its own log at +# logs/..log (created next to raw/, sources/, wiki/) with the +# full LLM request/response dump and tracebacks; the console keeps showing +# only the normal status lines and WARNING+ messages. Handy for unattended +# batch `add` runs where you only want to dig into the log after a failure. +# debug: true ``` | Key | Default | What it does | @@ -114,8 +122,25 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex | `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. | | `parallel_tool_calls` | unset | Whether the LLM agents (query, chat, lint, skill) may call tools in parallel. Unset keeps OpenKB's per-agent defaults; `true`/`false` force allow/sequential for every agent; `null` omits the setting (provider default). **Amazon Bedrock needs `null`** (see below). | | `entity_types` | 7 defaults | Custom vocabulary for entity pages. `other` is always kept. | +| `debug` | `false` | Enable per-file debug logging to `logs/.log` (see below) without needing `-v` on every command. | | `litellm:` | – | A pass-through block for LiteLLM. See below. | +### Debug logging + +`openkb -v ` and `debug: true` in `config.yaml` do the same thing: they +turn on OpenKB's DEBUG logging (full LLM request/response text, stage +tracebacks). The two differ only in scope and destination: + +- `-v` is a one-off, process-wide flag for any command. +- `debug: true` in `.openkb/config.yaml` is sticky per-KB, and only affects + `openkb add` (including a directory/batch add and the REST API's `/add`). + For each file being added, OpenKB writes a fresh `logs/..log` + (e.g. `logs/report.pdf.log`) next to `raw/`, `sources/`, `wiki/` — the + console still only shows the normal `Adding: ...` / `[OK]` / `[ERROR]` + status lines and any WARNING+ message, never the DEBUG detail. This keeps a + large batch `openkb add some-dir/` readable while still letting you dig + into exactly why an individual file failed afterward. + ### The `litellm:` block OpenKB forwards this block to LiteLLM so you can tune anything LiteLLM supports — diff --git a/openkb/cli.py b/openkb/cli.py index c9e543181..11b60a4d0 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -10,6 +10,7 @@ warnings.filterwarnings("ignore") import asyncio +import contextlib import json import logging from dataclasses import dataclass @@ -88,6 +89,71 @@ def filter(self, record: logging.LogRecord) -> bool: logger = logging.getLogger(__name__) +# Name of the per-KB directory that holds per-file debug logs (see +# ``_per_file_debug_log``), a sibling of raw/, sources/, wiki/. +DEBUG_LOGS_DIRNAME = "logs" + + +def _configure_console_logging() -> None: + """Set up root logging so the console only ever shows WARNING+ messages. + + Historically ``-v``/``--verbose`` raised the ``openkb`` logger to DEBUG + *and* let that flow straight to the console via the root handler + ``logging.basicConfig`` installs, which is unusable for a batch ``add`` + over many files (the request/response dump for every LLM call would + scroll past). DEBUG detail is instead captured per file by + ``_per_file_debug_log`` into ``logs/.log`` — this function just + pins the console handler's own level to WARNING so it can never emit + DEBUG/INFO records regardless of what level a logger further up allows + through (``-v`` still controls whether those records are created at all). + """ + root_logger = logging.getLogger() + if not root_logger.handlers: + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter("%(name)s %(levelname)s: %(message)s")) + root_logger.addHandler(handler) + for existing_handler in root_logger.handlers: + existing_handler.setLevel(logging.WARNING) + root_logger.setLevel(logging.WARNING) + + +@contextlib.contextmanager +def _per_file_debug_log(kb_dir: Path, file_path: Path, *, enabled: bool): + """Redirect DEBUG logging for one ``add``-pipeline file to its own log file. + + When ``enabled`` (from ``-v``/``--verbose`` or the KB's ``debug: true`` + config key — see ``_add_single_file_locked``), every ``openkb.*`` log + record produced while this context is active (conversion, PageIndex + indexing, LLM compilation) is additionally written to + ``kb_dir/logs/.log``, overwriting any log from a previous + run of the same file. The console is untouched: ``_configure_console_logging`` + already caps its handler at WARNING, so this is purely additive. A no-op + (yields immediately) when ``enabled`` is falsy, so callers can use this + unconditionally without branching. + """ + if not enabled: + yield + return + + logs_dir = kb_dir / DEBUG_LOGS_DIRNAME + logs_dir.mkdir(parents=True, exist_ok=True) + log_path = logs_dir / f"{file_path.name}.log" + + openkb_logger = logging.getLogger("openkb") + previous_level = openkb_logger.level + openkb_logger.setLevel(logging.DEBUG) + + file_handler = logging.FileHandler(log_path, mode="w", encoding="utf-8") + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s")) + openkb_logger.addHandler(file_handler) + try: + yield + finally: + openkb_logger.removeHandler(file_handler) + file_handler.close() + openkb_logger.setLevel(previous_level) + _KNOWN_PROVIDER_KEYS = ( "OPENAI_API_KEY", @@ -454,6 +520,31 @@ def add_single_file( def _add_single_file_locked( file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None +) -> Literal["added", "skipped", "failed"]: + """Resolve per-file debug logging, then run the add pipeline. + + Debug logging is enabled by ``-v``/``--verbose`` (process-wide, checked + via the ``openkb`` logger's effective level) or the KB's ``debug: true`` + config key; either source routes this file's DEBUG detail to + ``logs/.log`` instead of the console — see + ``_per_file_debug_log``. The actual conversion/index/compile pipeline is + in ``_add_single_file``. + """ + config = resolve_effective_config(kb_dir)[0] + debug_enabled = bool(config.get("debug")) or logging.getLogger("openkb").isEnabledFor( + logging.DEBUG + ) + with _per_file_debug_log(kb_dir, file_path, enabled=debug_enabled): + return _add_single_file(file_path, kb_dir, config, stage=stage, bundle=bundle) + + +def _add_single_file( + file_path: Path, + kb_dir: Path, + config: dict, + *, + stage: bool = True, + bundle=None, ) -> Literal["added", "skipped", "failed"]: """Convert, index, and compile a single document into the knowledge base. @@ -475,7 +566,6 @@ def _add_single_file_locked( from openkb.state import HashRegistry openkb_dir = kb_dir / ".openkb" - config = resolve_effective_config(kb_dir)[0] # The REST API passes a per-KB credential bundle so it never pollutes # process-wide state; only the CLI path needs the legacy global setup. if bundle is None: @@ -838,10 +928,7 @@ def append_cloud_log() -> None: @click.pass_context def cli(ctx, verbose, kb_dir_override): """OpenKB — Karpathy's LLM Knowledge Base workflow, powered by PageIndex.""" - logging.basicConfig( - format="%(name)s %(levelname)s: %(message)s", - level=logging.WARNING, - ) + _configure_console_logging() if verbose: logging.getLogger("openkb").setLevel(logging.DEBUG) ctx.ensure_object(dict) diff --git a/openkb/config.py b/openkb/config.py index 95ca8691f..3f9739913 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -36,6 +36,14 @@ # global/KB list overrides it wholesale; resolve_entity_types cleans the # effective value on read. "entity_types": list(DEFAULT_ENTITY_TYPES), + # Enables per-file DEBUG logging (redirected to logs/.log, never the + # console — see `openkb.cli._per_file_debug_log`) without needing the CLI's + # `-v`/`--verbose` flag on every invocation. Handy for unattended/batch + # `add` runs where a few files fail and the detailed traceback/LLM request + # dump is only needed after the fact. Deliberately NOT in GLOBAL_SCALAR_KEYS + # (KB config.yaml only, no global.yaml default / REST API exposure) — it's + # a CLI power-user diagnostic switch, not a workbench-editable setting. + "debug": False, } GLOBAL_CONFIG_DIR = Path.home() / ".config" / "openkb" diff --git a/tests/test_add_command.py b/tests/test_add_command.py index 3f51788e5..d37c6ef0b 100644 --- a/tests/test_add_command.py +++ b/tests/test_add_command.py @@ -3,8 +3,10 @@ from __future__ import annotations import json +import logging from unittest.mock import patch +import pytest from click.testing import CliRunner from openkb.cli import SUPPORTED_EXTENSIONS, _find_kb_dir, cli @@ -456,6 +458,133 @@ def test_add_requires_path_or_cloud(self, tmp_path): assert "Provide a PATH" in result.output +class TestPerFileDebugLogging: + """`debug: true` (config.yaml) / `-v` must route DEBUG detail to + logs/.log per processed file, and never to the console — see + `openkb.cli._configure_console_logging` / `_per_file_debug_log`.""" + + def _setup_kb(self, tmp_path): + (tmp_path / "raw").mkdir() + (tmp_path / "wiki" / "sources" / "images").mkdir(parents=True) + (tmp_path / "wiki" / "summaries").mkdir(parents=True) + (tmp_path / "wiki" / "concepts").mkdir(parents=True) + (tmp_path / "wiki" / "reports").mkdir(parents=True) + openkb_dir = tmp_path / ".openkb" + openkb_dir.mkdir() + (openkb_dir / "hashes.json").write_text(json.dumps({})) + return tmp_path + + @pytest.fixture(autouse=True) + def _reset_openkb_logger(self): + """Restore the shared `openkb` logger's level/handlers after each test + so a leaked DEBUG level or FileHandler never bleeds into other tests.""" + openkb_logger = logging.getLogger("openkb") + original_level = openkb_logger.level + original_handlers = list(openkb_logger.handlers) + yield + openkb_logger.setLevel(original_level) + for handler in list(openkb_logger.handlers): + if handler not in original_handlers: + openkb_logger.removeHandler(handler) + handler.close() + + def test_disabled_by_default_creates_no_logs_dir(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\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), + patch("openkb.cli._setup_llm_key"), + ): + outcome = add_single_file(doc, kb_dir) + + assert outcome == "added" + assert not (kb_dir / "logs").exists() + + def test_debug_config_writes_per_file_log(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\ndebug: true\n", encoding="utf-8" + ) + doc = tmp_path / "notes.md" + doc.write_text("# Notes\n\nBody", encoding="utf-8") + + async def fake_compile(*args, **kwargs): + logging.getLogger("openkb.agent.compiler").debug("marker: fake LLM request dump") + + with ( + patch( + "openkb.agent.compiler.compile_short_doc", + new_callable=AsyncMock, + side_effect=fake_compile, + ), + patch("openkb.cli._setup_llm_key"), + ): + outcome = add_single_file(doc, kb_dir) + + assert outcome == "added" + log_path = kb_dir / "logs" / "notes.md.log" + assert log_path.exists() + assert "marker: fake LLM request dump" in log_path.read_text(encoding="utf-8") + # `_per_file_debug_log` must restore the `openkb` logger's prior level + # once the file finishes, so a later non-debug file isn't affected. + assert logging.getLogger("openkb").level != logging.DEBUG + + def test_verbose_flag_also_triggers_per_file_log(self, tmp_path): + """`-v` (process-wide DEBUG on the `openkb` logger) must be honored + the same way as `debug: true`, even without setting it in config.yaml.""" + 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\n", encoding="utf-8") + doc = tmp_path / "report.md" + doc.write_text("# Report\n\nBody", encoding="utf-8") + + logging.getLogger("openkb").setLevel(logging.DEBUG) # simulates `-v` + + async def fake_compile(*args, **kwargs): + logging.getLogger("openkb.agent.compiler").debug("marker: verbose path") + + with ( + patch( + "openkb.agent.compiler.compile_short_doc", + new_callable=AsyncMock, + side_effect=fake_compile, + ), + patch("openkb.cli._setup_llm_key"), + ): + outcome = add_single_file(doc, kb_dir) + + assert outcome == "added" + log_path = kb_dir / "logs" / "report.md.log" + assert log_path.exists() + assert "marker: verbose path" in log_path.read_text(encoding="utf-8") + + def test_configure_console_logging_caps_handler_at_warning(self): + """Even when a logger's own level allows DEBUG through (e.g. after + `-v`), the console handler itself must never emit below WARNING — + DEBUG detail is only ever written to a per-file log (see + `_per_file_debug_log`).""" + from openkb.cli import _configure_console_logging + + _configure_console_logging() + root_handlers = logging.getLogger().handlers + assert root_handlers + assert all(h.level >= logging.WARNING for h in root_handlers) + + class TestImportFromPageindexCloud: def _setup_kb(self, tmp_path): (tmp_path / "raw").mkdir() From 6c613b0d1baddae4f1dacb3de6b18f2caaf98369 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Mon, 31 Aug 2026 08:07:33 +0200 Subject: [PATCH 13/16] fix(agent): log stream chunk phase as start/end, not one line per chunk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-chunk DEBUG logging (_log_chunk_timing) drowned out the rest of a log on a long response — hundreds of LLM stream chunk [...] #N lines for a single LLM call, e.g. one concepts-plan request logging 205 individual chunk lines. Replaces it with exactly two log lines per LLM call: - _log_stream_start: logged once, when the first chunk arrives (time-to-first-token). - _log_stream_end: logged once, when the stream finishes cleanly (total chunk count + elapsed time for the last chunk). - _log_stream_interrupted: logged once instead of _log_stream_end if the stream raises mid-iteration — reports how many chunks were successfully received and when, right before the exception is re-raised (still a complete failure, no partial buffer). Special-cases zero chunks (failure before any byte arrived) with dedicated wording instead of an inapplicable chunk number. Updated tests/test_compiler.py::TestLLMStreamTimingDebugLogging to match: asserts exactly one start + one end/interrupted line, and the explicit absence of the old per-chunk lines. --- openkb/agent/compiler.py | 105 +++++++++++++++++++++++------------- tests/test_compiler.py | 112 +++++++++++++++++++++++---------------- 2 files changed, 135 insertions(+), 82 deletions(-) diff --git a/openkb/agent/compiler.py b/openkb/agent/compiler.py index 311ce4f16..756f66ec7 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -414,33 +414,72 @@ def _merge_stream_chunks(chunks: list, messages: list[dict]): return litellm.stream_chunk_builder(chunks, messages=messages) -def _log_chunk_timing( - step_name: str, chunk_number: int, t0: float, last_t: float, now: float +def _log_stream_start(step_name: str, t0: float, first_chunk_t: float) -> None: + """Debug-log the time-to-first-chunk (TTFT) once a stream's first chunk arrives. + + Marks the start of a "chunk phase" in the log. The counterpart is + :func:`_log_stream_end` (clean finish) or :func:`_log_stream_interrupted` + (mid-stream failure) — together these replace a debug line per chunk + (which used to drown out the rest of the log on a long response, e.g. + hundreds of lines for one LLM call) with exactly one line at the start + and exactly one more at the end/interruption. + """ + logger.debug( + "LLM stream started [%s]: first chunk after %.2fs", + step_name, + first_chunk_t - t0, + ) + + +def _log_stream_end(step_name: str, chunk_count: int, t0: float, last_chunk_t: float) -> None: + """Debug-log a stream's clean completion: total chunk count and elapsed time.""" + logger.debug( + "LLM stream finished [%s]: %d chunk(s), last chunk after %.2fs total", + step_name, + chunk_count, + last_chunk_t - t0, + ) + + +def _log_stream_interrupted( + step_name: str, chunk_count: int, t0: float, last_chunk_t: float ) -> None: - """Debug-log arrival timing for one streamed chunk. - - Run with ``openkb -v`` to see, per LLM call, how long the first chunk took - (time-to-first-token) and the gap to each subsequent chunk. If chunks do - arrive before a timeout, that shows the proxy/gateway is forwarding the - stream; if no chunk is ever logged before a timeout, the instrumentation - narrows the problem to a no-first-byte case (e.g. slow TTFT or an - intermediary buffering the response). + """Debug-log a stream that raised mid-iteration, right before it is re-raised. + + ``chunk_count`` is how many chunks were successfully received before the + failure (0 if the very first chunk never arrived). The exception itself + (with traceback) is attached via ``exc_info=True`` so the failure and the + chunk-phase summary land in a single log record. """ + now = time.time() + if chunk_count == 0: + logger.debug( + "LLM stream [%s] interrupted unexpectedly before any chunk arrived (%.2fs total)", + step_name, + now - t0, + exc_info=True, + ) + return logger.debug( - "LLM stream chunk [%s] #%d after %.2fs total (+%.2fs since previous)", + "LLM stream [%s] interrupted unexpectedly after chunk %d " + "(last chunk after %.2fs, failure after %.2fs total)", step_name, - chunk_number, + chunk_count, + last_chunk_t - t0, now - t0, - now - last_t, + exc_info=True, ) def _consume_stream(stream, step_name: str, t0: float) -> list: - """Collect a sync LiteLLM stream into a list, logging per-chunk timing. - - A mid-stream exception (e.g. the gateway idle-timeout firing) propagates - after logging how many chunks arrived and when, so callers still see a - complete failure — no partial buffer is ever returned. + """Collect a sync LiteLLM stream into a list, debug-logging the chunk phase. + + Logs exactly one line when the first chunk arrives (time-to-first-token) + and exactly one more line when the stream ends — either + :func:`_log_stream_end` on a clean finish or :func:`_log_stream_interrupted` + if it raises mid-iteration. A mid-stream exception (e.g. the gateway + idle-timeout firing) propagates after being logged, so callers still see + a complete failure — no partial buffer is ever returned. """ if not logger.isEnabledFor(logging.DEBUG): return list(stream) @@ -450,26 +489,22 @@ def _consume_stream(stream, step_name: str, t0: float) -> list: try: for chunk in stream: now = time.time() - _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now) + if not chunks: + _log_stream_start(step_name, t0, now) chunks.append(chunk) last_t = now except Exception: - logger.debug( - "LLM stream [%s] failed after %.2fs with %d chunk(s) received", - step_name, - time.time() - t0, - len(chunks), - exc_info=True, - ) + _log_stream_interrupted(step_name, len(chunks), t0, last_t) raise + _log_stream_end(step_name, len(chunks), t0, last_t) return chunks async def _consume_stream_async(stream, step_name: str, t0: float) -> list: - """Collect an async LiteLLM stream into a list, logging per-chunk timing. + """Collect an async LiteLLM stream into a list, debug-logging the chunk phase. - Mirrors :func:`_consume_stream`, including the no-partial-buffer invariant - on failure. + Mirrors :func:`_consume_stream`, including the start/end-or-interrupted + logging and the no-partial-buffer invariant on failure. """ if not logger.isEnabledFor(logging.DEBUG): return [chunk async for chunk in stream] @@ -479,18 +514,14 @@ async def _consume_stream_async(stream, step_name: str, t0: float) -> list: try: async for chunk in stream: now = time.time() - _log_chunk_timing(step_name, len(chunks) + 1, t0, last_t, now) + if not chunks: + _log_stream_start(step_name, t0, now) chunks.append(chunk) last_t = now except Exception: - logger.debug( - "LLM stream [%s] failed after %.2fs with %d chunk(s) received", - step_name, - time.time() - t0, - len(chunks), - exc_info=True, - ) + _log_stream_interrupted(step_name, len(chunks), t0, last_t) raise + _log_stream_end(step_name, len(chunks), t0, last_t) return chunks diff --git a/tests/test_compiler.py b/tests/test_compiler.py index a5119ff2a..34a94ac6c 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -2710,9 +2710,11 @@ async def test_llm_call_async_injects_extra_headers(self): class TestLLMStreamTimingDebugLogging: - """Per-chunk debug timing should be visible when verbose logging is enabled.""" + """Chunk-phase debug logging should be visible when verbose logging is + enabled: one line when the first chunk arrives, one more when the stream + ends cleanly or is interrupted — never one line per chunk (see #).""" - def test_llm_call_logs_each_stream_chunk_at_debug(self, caplog): + def test_llm_call_logs_stream_start_and_end_at_debug(self, caplog): from openkb.agent.compiler import _llm_call caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") @@ -2728,17 +2730,17 @@ def test_llm_call_logs_each_stream_chunk_at_debug(self, caplog): out = _llm_call("m", [{"role": "user", "content": "hi"}], "sync-step") assert out == "ok" - chunk_logs = [ - record.getMessage() - for record in caplog.records - if "LLM stream chunk [sync-step]" in record.getMessage() - ] - assert len(chunk_logs) == 3 - assert "#1" in chunk_logs[0] - assert "#3" in chunk_logs[-1] + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [sync-step]" in m] + end_logs = [m for m in messages if "LLM stream finished [sync-step]" in m] + # Exactly one start line and one end line — never a line per chunk. + assert len(start_logs) == 1 + assert len(end_logs) == 1 + assert "3 chunk(s)" in end_logs[0] + assert not any("LLM stream chunk [sync-step]" in m for m in messages) @pytest.mark.asyncio - async def test_llm_call_async_logs_each_stream_chunk_at_debug(self, caplog): + async def test_llm_call_async_logs_stream_start_and_end_at_debug(self, caplog): from openkb.agent.compiler import _llm_call_async caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") @@ -2751,16 +2753,15 @@ async def test_llm_call_async_logs_each_stream_chunk_at_debug(self, caplog): out = await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-step") assert out == "ok" - chunk_logs = [ - record.getMessage() - for record in caplog.records - if "LLM stream chunk [async-step]" in record.getMessage() - ] - assert len(chunk_logs) == 3 - assert "#1" in chunk_logs[0] - assert "#3" in chunk_logs[-1] - - def test_llm_call_logs_stream_failure_before_reraising(self, caplog): + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [async-step]" in m] + end_logs = [m for m in messages if "LLM stream finished [async-step]" in m] + assert len(start_logs) == 1 + assert len(end_logs) == 1 + assert "3 chunk(s)" in end_logs[0] + assert not any("LLM stream chunk [async-step]" in m for m in messages) + + def test_llm_call_logs_interruption_before_reraising(self, caplog): from openkb.agent.compiler import _llm_call caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") @@ -2776,22 +2777,21 @@ def test_llm_call_logs_stream_failure_before_reraising(self, caplog): with pytest.raises(RuntimeError, match="stream exploded"): _llm_call("m", [{"role": "user", "content": "hi"}], "sync-fail-step") - chunk_logs = [ - record.getMessage() - for record in caplog.records - if "LLM stream chunk [sync-fail-step]" in record.getMessage() - ] - failure_logs = [ - record.getMessage() - for record in caplog.records - if "LLM stream [sync-fail-step] failed after" in record.getMessage() + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [sync-fail-step]" in m] + interrupted_logs = [ + m for m in messages if "LLM stream [sync-fail-step] interrupted unexpectedly" in m ] - assert len(chunk_logs) == 2 - assert len(failure_logs) == 1 - assert "2 chunk(s) received" in failure_logs[0] + # Exactly one start line and one interruption line, no end-of-stream line, + # and no per-chunk lines in between. + assert len(start_logs) == 1 + assert len(interrupted_logs) == 1 + assert "after chunk 2" in interrupted_logs[0] + assert not any("LLM stream finished [sync-fail-step]" in m for m in messages) + assert not any("LLM stream chunk [sync-fail-step]" in m for m in messages) @pytest.mark.asyncio - async def test_llm_call_async_logs_stream_failure_before_reraising(self, caplog): + async def test_llm_call_async_logs_interruption_before_reraising(self, caplog): from openkb.agent.compiler import _llm_call_async caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") @@ -2808,19 +2808,41 @@ async def test_llm_call_async_logs_stream_failure_before_reraising(self, caplog) with pytest.raises(RuntimeError, match="stream exploded"): await _llm_call_async("m", [{"role": "user", "content": "hi"}], "async-fail-step") - chunk_logs = [ - record.getMessage() - for record in caplog.records - if "LLM stream chunk [async-fail-step]" in record.getMessage() + messages = [record.getMessage() for record in caplog.records] + start_logs = [m for m in messages if "LLM stream started [async-fail-step]" in m] + interrupted_logs = [ + m for m in messages if "LLM stream [async-fail-step] interrupted unexpectedly" in m ] - failure_logs = [ - record.getMessage() - for record in caplog.records - if "LLM stream [async-fail-step] failed after" in record.getMessage() + assert len(start_logs) == 1 + assert len(interrupted_logs) == 1 + assert "after chunk 2" in interrupted_logs[0] + assert not any("LLM stream finished [async-fail-step]" in m for m in messages) + assert not any("LLM stream chunk [async-fail-step]" in m for m in messages) + + def test_llm_call_logs_interruption_before_any_chunk(self, caplog): + """No chunk ever arrives (e.g. a proxy silently buffering despite + stream=True): no start line, and the interruption line says so + instead of an inapplicable chunk number.""" + from openkb.agent.compiler import _llm_call + + caplog.set_level(logging.DEBUG, logger="openkb.agent.compiler") + error = RuntimeError("connect timeout") + with ( + patch("openkb.agent.compiler._Spinner", _NoOpSpinner), + patch("openkb.agent.compiler.litellm") as mock_litellm, + ): + mock_litellm.completion = MagicMock(return_value=_stream_then_raise([], error)) + + with pytest.raises(RuntimeError, match="connect timeout"): + _llm_call("m", [{"role": "user", "content": "hi"}], "sync-no-chunk-step") + + messages = [record.getMessage() for record in caplog.records] + interrupted_logs = [ + m for m in messages if "LLM stream [sync-no-chunk-step] interrupted unexpectedly" in m ] - assert len(chunk_logs) == 2 - assert len(failure_logs) == 1 - assert "2 chunk(s) received" in failure_logs[0] + assert len(interrupted_logs) == 1 + assert "before any chunk arrived" in interrupted_logs[0] + assert not any("LLM stream started [sync-no-chunk-step]" in m for m in messages) class TestCacheControlStripping: From c102823a305dbe0755b629885385f5142f0e064e Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Mon, 31 Aug 2026 08:08:32 +0200 Subject: [PATCH 14/16] fix(cli): delete per-file debug log only on an exact "added" outcome MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the log written by _per_file_debug_log was always kept regardless of outcome — for a batch add that succeeds on most files and only fails on a few, this left a debug log for every single file, defeating the point of routing DEBUG detail out of the console (still noisy, just moved to disk). _per_file_debug_log now yields the log file's Path (or None when disabled) instead of nothing, so _add_single_file_locked can inspect it once the outcome is known and the file handler has already been closed. The log is deleted ONLY on an exact outcome == "added" match — never != "failed" — so a future third outcome (e.g. a partial/ degraded success) doesn't silently inherit delete-on-success just because it isn't "failed"; it would have to be added to the delete condition explicitly. "skipped" (dedup — nothing new happened) and "failed" (the whole point of turning this on) both keep their log. Updated tests/test_add_command.py::TestPerFileDebugLogging: - test_debug_config_writes_per_file_log_deleted_on_added (renamed from test_debug_config_writes_per_file_log): now asserts deletion. - test_debug_log_retained_on_failed (new). - test_debug_log_retained_on_skipped (new). - test_verbose_flag_also_triggers_per_file_log: switched to a forced failure so the log's content can still be asserted post-run. --- examples/configuration/README.md | 6 ++- openkb/cli.py | 27 +++++++++--- tests/test_add_command.py | 74 ++++++++++++++++++++++++-------- 3 files changed, 83 insertions(+), 24 deletions(-) diff --git a/examples/configuration/README.md b/examples/configuration/README.md index d9d889254..f450a37b3 100644 --- a/examples/configuration/README.md +++ b/examples/configuration/README.md @@ -139,7 +139,11 @@ tracebacks). The two differ only in scope and destination: console still only shows the normal `Adding: ...` / `[OK]` / `[ERROR]` status lines and any WARNING+ message, never the DEBUG detail. This keeps a large batch `openkb add some-dir/` readable while still letting you dig - into exactly why an individual file failed afterward. + into exactly why an individual file failed afterward. The log is deleted + automatically once the file's outcome is exactly `"added"` (a fresh, + successful compile has nothing left to debug); it is kept on `"skipped"` + (dedup — nothing new happened) and `"failed"` (the whole point of turning + this on) so you can inspect it after the batch finishes. ### The `litellm:` block diff --git a/openkb/cli.py b/openkb/cli.py index 11b60a4d0..e9b573920 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -128,11 +128,16 @@ def _per_file_debug_log(kb_dir: Path, file_path: Path, *, enabled: bool): ``kb_dir/logs/.log``, overwriting any log from a previous run of the same file. The console is untouched: ``_configure_console_logging`` already caps its handler at WARNING, so this is purely additive. A no-op - (yields immediately) when ``enabled`` is falsy, so callers can use this - unconditionally without branching. + (yields ``None`` immediately) when ``enabled`` is falsy, so callers can use + this unconditionally without branching. + + Yields the log file's path (or ``None`` when disabled) so the caller can + decide, once the outcome of the add is known and this context has already + closed the file handler, whether to keep or delete it — see + ``_add_single_file_locked``, which deletes it only on ``"added"``. """ if not enabled: - yield + yield None return logs_dir = kb_dir / DEBUG_LOGS_DIRNAME @@ -148,7 +153,7 @@ def _per_file_debug_log(kb_dir: Path, file_path: Path, *, enabled: bool): file_handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s")) openkb_logger.addHandler(file_handler) try: - yield + yield log_path finally: openkb_logger.removeHandler(file_handler) file_handler.close() @@ -529,13 +534,23 @@ def _add_single_file_locked( ``logs/.log`` instead of the console — see ``_per_file_debug_log``. The actual conversion/index/compile pipeline is in ``_add_single_file``. + + The log is deleted once the outcome is known, but ONLY on an exact + ``"added"`` match — not e.g. ``!= "failed"`` — so a future third outcome + (say, a partial/degraded success) can be added later without silently + inheriting delete-on-success just because it isn't "failed". "skipped" + keeps its log too: a dedup hit is not a fresh compile run, so there is + nothing new in it to discard. """ config = resolve_effective_config(kb_dir)[0] debug_enabled = bool(config.get("debug")) or logging.getLogger("openkb").isEnabledFor( logging.DEBUG ) - with _per_file_debug_log(kb_dir, file_path, enabled=debug_enabled): - return _add_single_file(file_path, kb_dir, config, stage=stage, bundle=bundle) + with _per_file_debug_log(kb_dir, file_path, enabled=debug_enabled) as log_path: + outcome = _add_single_file(file_path, kb_dir, config, stage=stage, bundle=bundle) + if log_path is not None and outcome == "added": + log_path.unlink(missing_ok=True) + return outcome def _add_single_file( diff --git a/tests/test_add_command.py b/tests/test_add_command.py index d37c6ef0b..d3d61cd2e 100644 --- a/tests/test_add_command.py +++ b/tests/test_add_command.py @@ -507,7 +507,9 @@ def test_disabled_by_default_creates_no_logs_dir(self, tmp_path): assert outcome == "added" assert not (kb_dir / "logs").exists() - def test_debug_config_writes_per_file_log(self, tmp_path): + def test_debug_config_writes_per_file_log_deleted_on_added(self, tmp_path): + """A successful ("added") run's log is deleted once the outcome is + known — a fresh, successful compile has nothing left to debug.""" from unittest.mock import AsyncMock from openkb.cli import add_single_file @@ -533,18 +535,62 @@ async def fake_compile(*args, **kwargs): outcome = add_single_file(doc, kb_dir) assert outcome == "added" - log_path = kb_dir / "logs" / "notes.md.log" - assert log_path.exists() - assert "marker: fake LLM request dump" in log_path.read_text(encoding="utf-8") + assert not (kb_dir / "logs" / "notes.md.log").exists() # `_per_file_debug_log` must restore the `openkb` logger's prior level # once the file finishes, so a later non-debug file isn't affected. assert logging.getLogger("openkb").level != logging.DEBUG + def test_debug_log_retained_on_failed(self, tmp_path): + """A "failed" run keeps its log — that's the entire point of turning + debug logging on (inspecting exactly why a file failed).""" + 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\ndebug: true\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", side_effect=RuntimeError("boom")), + patch("openkb.cli.time.sleep"), + patch("openkb.cli._setup_llm_key"), + ): + outcome = add_single_file(doc, kb_dir) + + assert outcome == "failed" + log_path = kb_dir / "logs" / "notes.md.log" + assert log_path.exists() + assert "Compilation traceback" in log_path.read_text(encoding="utf-8") + + def test_debug_log_retained_on_skipped(self, tmp_path): + """A "skipped" (dedup) run keeps its log too — an explicit equality + check on "added" (not e.g. `!= "failed"`) is what makes this possible: + a later third outcome would default to "kept" unless added to the + delete condition explicitly.""" + from openkb.cli import add_single_file + from openkb.converter import ConvertResult + + kb_dir = self._setup_kb(tmp_path) + (kb_dir / ".openkb" / "config.yaml").write_text( + "model: gpt-4o-mini\ndebug: true\n", encoding="utf-8" + ) + doc = tmp_path / "notes.md" + doc.write_text("# Notes\n\nBody", encoding="utf-8") + + with patch("openkb.cli.convert_document", return_value=ConvertResult(skipped=True)): + outcome = add_single_file(doc, kb_dir) + + assert outcome == "skipped" + assert (kb_dir / "logs" / "notes.md.log").exists() + def test_verbose_flag_also_triggers_per_file_log(self, tmp_path): """`-v` (process-wide DEBUG on the `openkb` logger) must be honored - the same way as `debug: true`, even without setting it in config.yaml.""" - from unittest.mock import AsyncMock - + the same way as `debug: true`, even without setting it in config.yaml. + Uses a forced failure so the log is retained and its content can be + asserted (a success would delete it — see + `test_debug_config_writes_per_file_log_deleted_on_added`).""" from openkb.cli import add_single_file kb_dir = self._setup_kb(tmp_path) @@ -554,23 +600,17 @@ def test_verbose_flag_also_triggers_per_file_log(self, tmp_path): logging.getLogger("openkb").setLevel(logging.DEBUG) # simulates `-v` - async def fake_compile(*args, **kwargs): - logging.getLogger("openkb.agent.compiler").debug("marker: verbose path") - with ( - patch( - "openkb.agent.compiler.compile_short_doc", - new_callable=AsyncMock, - side_effect=fake_compile, - ), + patch("openkb.agent.compiler.compile_short_doc", side_effect=RuntimeError("boom")), + patch("openkb.cli.time.sleep"), patch("openkb.cli._setup_llm_key"), ): outcome = add_single_file(doc, kb_dir) - assert outcome == "added" + assert outcome == "failed" log_path = kb_dir / "logs" / "report.md.log" assert log_path.exists() - assert "marker: verbose path" in log_path.read_text(encoding="utf-8") + assert "Compilation traceback" in log_path.read_text(encoding="utf-8") def test_configure_console_logging_caps_handler_at_warning(self): """Even when a logger's own level allows DEBUG through (e.g. after From 16b9eb0cce60252e3f95f450376a097441e1a384 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Mon, 31 Aug 2026 10:23:49 +0200 Subject: [PATCH 15/16] feat(cli): configurable insert_mode (fail-fast / fail-at-end) for strict compile-failure handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add `insert_mode` config key (`.openkb/config.yaml`) with three values: - "normal" (default): unchanged behavior — a concept/entity generation failure during compile is logged as a warning and the file is still reported "added". - "fail-fast": the first concept/entity generation failure cancels every other still-pending generation in the batch and immediately raises `ConceptCompilationError` — nothing from the batch is written. - "fail-at-end": every planned concept/entity generation is attempted (so every failure for the document is logged in one pass) before `ConceptCompilationError` is raised if anything failed. - Both strict modes rely entirely on the existing mutation-snapshot rollback (`openkb.add_coordinator`/`openkb.mutation`) to discard the add and report it "failed" — no new rollback path needed. The existing "keep raw/ on failed" and "keep the debug log on a non-'added' outcome" behaviors already cover the raw-file and log-preservation requirements for strict mode. - `_compile_concepts`'s three early-return paths (unparseable plan, scalar plan, all-items-filtered-as-malformed) now also raise under a strict insert_mode, not just individual concept/entity generation failures — a genuinely empty plan (nothing was ever planned) still counts as complete success in every mode. - `compile_short_doc`/`compile_long_doc` resolve `insert_mode` from the already-loaded KB config, so no CLI-level plumbing is needed. Resolves #239 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- config.yaml.example | 10 ++ examples/configuration/README.md | 11 ++ openkb/agent/compiler.py | 125 ++++++++++++++++++-- openkb/config.py | 26 +++++ tests/test_compiler.py | 195 +++++++++++++++++++++++++++++++ tests/test_config.py | 19 +++ 6 files changed, 374 insertions(+), 12 deletions(-) diff --git a/config.yaml.example b/config.yaml.example index 47d2fda2b..324bd0ebe 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -8,6 +8,16 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex # large PDFs. Omit to let each stage apply its own default. # concurrency: 5 +# Optional: how a partially-failed compile (some concept/entity couldn't be +# generated, e.g. a transient LLM error) is reported for a single `add`. +# normal (default) log a warning, still report the file "added". +# fail-fast abort as soon as the first concept/entity fails; the add is +# rolled back and reported "failed". +# fail-at-end attempt every planned concept/entity first (so every failure +# for the file is logged in one pass), then roll back and +# report "failed" if anything failed. +# insert_mode: normal + # Optional: whether the LLM agents (query, chat, lint, skill) may call tools # in parallel. Leave it UNSET (commented out) to keep OpenKB's per-agent # defaults. Setting it applies the SAME value to every agent: diff --git a/examples/configuration/README.md b/examples/configuration/README.md index 25e5b3a85..6009dbfb7 100644 --- a/examples/configuration/README.md +++ b/examples/configuration/README.md @@ -76,6 +76,16 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex # large PDFs. Omit to let each stage apply its own default. # concurrency: 5 +# Optional: how a partially-failed compile (some concept/entity couldn't be +# generated, e.g. a transient LLM error) is reported for a single `add`. +# normal (default) log a warning, still report the file "added". +# fail-fast abort as soon as the first concept/entity fails; the add is +# rolled back and reported "failed". +# fail-at-end attempt every planned concept/entity first (so every failure +# for the file is logged in one pass), then roll back and +# report "failed" if anything failed. +# insert_mode: normal + # Optional: whether the LLM agents (query, chat, lint, skill) may call tools # in parallel. Leave it UNSET (commented out) to keep OpenKB's per-agent # defaults. Setting it applies the SAME value to every agent: @@ -112,6 +122,7 @@ pageindex_threshold: 20 # PDF pages threshold for PageIndex | `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. | +| `insert_mode` | `normal` | How a partially-failed compile (some concept/entity couldn't be generated) is reported. `normal` logs a warning and still reports the file "added". `fail-fast` aborts on the first failure; `fail-at-end` attempts every planned concept/entity first (so every failure is logged in one pass). Both strict modes roll back the add and report it "failed" instead of "added". | | `parallel_tool_calls` | unset | Whether the LLM agents (query, chat, lint, skill) may call tools in parallel. Unset keeps OpenKB's per-agent defaults; `true`/`false` force allow/sequential for every agent; `null` omits the setting (provider default). **Amazon Bedrock needs `null`** (see below). | | `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/agent/compiler.py b/openkb/agent/compiler.py index d0c9f878d..84ae52727 100644 --- a/openkb/agent/compiler.py +++ b/openkb/agent/compiler.py @@ -397,6 +397,19 @@ class TruncatedResponseError(Exception): treat truncation as a failure (so a partial page is skipped, not written).""" +class ConceptCompilationError(Exception): + """Raised by ``_compile_concepts`` when ``insert_mode`` is ``"fail-fast"`` + or ``"fail-at-end"`` and one or more planned concept/entity updates could + not be generated for a document (see ``openkb.config.resolve_insert_mode``). + + Propagates through ``compile_short_doc``/``compile_long_doc`` up to + ``cli._add_single_file_locked``'s ``commit_body``, where the existing + mutation-snapshot rollback (``openkb.add_coordinator``) already reverts + every wiki/raw change for the add and reports the file as ``"failed"`` — + no separate rollback path is needed for strict mode. + """ + + def _llm_call( model: str, messages: list[dict], @@ -1604,6 +1617,7 @@ async def _compile_concepts( rewrite_summary: bool = False, entity_types: list[str] | None = None, bundle=None, + insert_mode: str = "normal", ) -> None: """Shared Steps 2-4: concepts plan → generate/update → index. @@ -1613,6 +1627,27 @@ async def _compile_concepts( written to disk. When ``rewrite_summary=True`` (short-doc path), the summary is rewritten by the LLM after concepts are finalized so its wikilinks reflect the actual concept pages on disk. + + ``insert_mode`` (see ``openkb.config.resolve_insert_mode``) controls what + happens when one or more planned concept/entity updates cannot be + generated: + + - ``"normal"`` (default): unchanged behavior — failures are logged as + warnings and whatever *did* generate is written; the document as a + whole is still considered compiled. + - ``"fail-fast"``: the first concept/entity generation failure cancels + every other still-pending (not yet started) generation in this batch + and immediately raises ``ConceptCompilationError`` — nothing from this + batch is written. + - ``"fail-at-end"``: every planned concept/entity generation is attempted + (so every failure for this document is logged in one pass) and + whatever succeeded is written, same as "normal" — but + ``ConceptCompilationError`` is raised at the end if anything failed. + + In both strict modes the raised exception is expected to propagate out of + ``compile_short_doc``/``compile_long_doc`` so the caller's existing + mutation rollback discards this add entirely (see + ``ConceptCompilationError``). """ source_file = f"summaries/{doc_name}.md" @@ -1627,6 +1662,20 @@ async def _compile_concepts( concept_briefs = _read_concept_briefs(wiki_dir) entity_briefs = _read_entity_briefs(wiki_dir) + def _maybe_raise_incomplete(reason: str) -> None: + """Raise ``ConceptCompilationError`` under a strict ``insert_mode``. + + A no-op under ``"normal"``, matching today's silent-partial-success + behavior. Called from every early-return branch below plus the final + completeness check, so both strict modes cover the "plan came back + unparseable/empty" cases, not just individual concept/entity + generation failures. + """ + if insert_mode in ("fail-fast", "fail-at-end"): + raise ConceptCompilationError( + f"insert_mode={insert_mode!r}: {doc_name!r} compiled incompletely — {reason}" + ) + # Second cache breakpoint: end of the assistant summary message. Covers # (system + doc + summary) for the plan call and every concept call. summary_msg = {"role": "assistant", "content": _cached_text(summary)} @@ -1686,6 +1735,7 @@ def _write_v1_summary_stripped() -> None: f"no concept pages generated. See log (stderr) for details.\n" ) sys.stdout.flush() + _maybe_raise_incomplete("concepts plan response was unparseable") if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) @@ -1706,6 +1756,7 @@ def _write_v1_summary_stripped() -> None: type(parsed).__name__, doc_name, ) + _maybe_raise_incomplete("concepts plan parsed to a scalar, not a usable plan") if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) @@ -1788,6 +1839,13 @@ def _raw_group_count(group: object) -> int: and not entity_update and not entity_related ): + # A genuinely empty plan (original_total == 0) is a complete, valid + # outcome for strict modes too — nothing was planned, so nothing is + # missing. But if items were planned and all got dropped as malformed + # (original_total > 0, already warned above), that's real content + # loss under a strict insert_mode. + if original_total > 0: + _maybe_raise_incomplete("all planned concept/entity items were dropped as malformed") if rewrite_summary: _write_v1_summary_stripped() _update_index(wiki_dir, doc_name, [], doc_brief=doc_brief, doc_type=doc_type) @@ -1964,16 +2022,18 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: return name, content, brief, etype_out tasks = [] - tasks.extend(_gen_create(c) for c in create_items) - tasks.extend(_gen_update(c) for c in update_items) + tasks.extend(asyncio.create_task(_gen_create(c)) for c in create_items) + tasks.extend(asyncio.create_task(_gen_update(c)) for c in update_items) # --- Step 3 (entities): build the entity task list up front so it can be # gathered concurrently with the concept tasks below. Entity coroutines # return 4-arity tuples (name, content, brief, type), so their results are # processed in their own loop rather than mixed with the concept tuples. + # Wrapped in asyncio.create_task (not left as bare coroutines) so a + # "fail-fast" insert_mode can cancel the ones still pending below. entity_tasks = [] - entity_tasks.extend(_gen_entity_create(e) for e in entity_create) - entity_tasks.extend(_gen_entity_update(e) for e in entity_update) + entity_tasks.extend(asyncio.create_task(_gen_entity_create(e)) for e in entity_create) + entity_tasks.extend(asyncio.create_task(_gen_entity_update(e)) for e in entity_update) concept_names: list[str] = [] concept_briefs_map: dict[str, str] = {} @@ -1998,13 +2058,37 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: results, entity_results = ([], []) if tasks or entity_tasks: - results, entity_results = await asyncio.gather( - asyncio.gather(*tasks, return_exceptions=True), - asyncio.gather(*entity_tasks, return_exceptions=True), - ) + if insert_mode == "fail-fast": + # Wait only until the first exception surfaces (or everything + # finishes cleanly) instead of always waiting for the full batch — + # cancelling whatever hasn't started/finished yet saves the LLM + # calls that batch would have made. asyncio.wait requires Tasks + # (not bare coroutines), hence the create_task() wrapping above. + all_tasks = tasks + entity_tasks + done, pending = await asyncio.wait(all_tasks, return_when=asyncio.FIRST_EXCEPTION) + first_exc = next((t.exception() for t in done if t.exception() is not None), None) + if first_exc is not None: + for t in pending: + t.cancel() + if pending: + # Swallow the resulting CancelledErrors; we only need the + # cancellations to settle before raising below. + await asyncio.gather(*pending, return_exceptions=True) + logger.warning("Concept/entity generation failed: %s", first_exc) + raise ConceptCompilationError( + f"insert_mode='fail-fast': aborting compile for {doc_name!r} after a " + f"concept/entity generation failure: {first_exc}" + ) from first_exc + results = [t.result() for t in tasks] + entity_results = [t.result() for t in entity_tasks] + else: + results, entity_results = await asyncio.gather( + asyncio.gather(*tasks, return_exceptions=True), + asyncio.gather(*entity_tasks, return_exceptions=True), + ) + failure_types: list[str] = [] if tasks: - failure_types: list[str] = [] for r in results: if isinstance(r, Exception): logger.warning("Concept generation failed: %s", r) @@ -2028,8 +2112,8 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: ) sys.stdout.flush() + entity_failure_types: list[str] = [] if entity_tasks: - entity_failure_types: list[str] = [] for r in entity_results: if isinstance(r, Exception): logger.warning("Entity generation failed: %s", r) @@ -2052,6 +2136,7 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: sys.stdout.flush() # Strip ghost wikilinks from entity bodies and write each page. + for name, page_content, brief, etype in entity_pending: cleaned, ghosts = strip_ghost_wikilinks(page_content, known_targets) if ghosts: @@ -2201,6 +2286,18 @@ async def _gen_entity_update(ent: dict) -> tuple[str, str, str, str]: entity_meta=entity_meta, ) + # "fail-fast" always raises earlier (see the gather branch above) before + # reaching this point, so only "fail-at-end" needs a completeness check + # here — everything planned was attempted (so every failure for this + # document is already logged above), and only now do we decide whether + # the document as a whole should count as failed. + if insert_mode == "fail-at-end" and (failure_types or entity_failure_types): + raise ConceptCompilationError( + f"insert_mode='fail-at-end': {doc_name!r} had {len(failure_types)} failed " + f"concept(s) and {len(entity_failure_types)} failed entity(ies) — " + f"{', '.join(sorted(set(failure_types + entity_failure_types))) or 'see log (stderr)'}" + ) + async def compile_short_doc( doc_name: str, @@ -2215,11 +2312,12 @@ async def compile_short_doc( Step 1: Build base context A (schema + doc content), generate summary. Steps 2-4: Delegated to ``_compile_concepts``. """ - from openkb.config import resolve_effective_config + from openkb.config import resolve_effective_config, resolve_insert_mode config = resolve_effective_config(kb_dir)[0] language: str = config.get("language", "en") entity_types = resolve_entity_types(config) + insert_mode = resolve_insert_mode(config) wiki_dir = kb_dir / "wiki" schema_md = get_agents_md(wiki_dir) @@ -2281,6 +2379,7 @@ async def compile_short_doc( rewrite_summary=True, entity_types=entity_types, bundle=bundle, + insert_mode=insert_mode, ) finally: # Close per-loop litellm async clients before asyncio.run tears this @@ -2303,11 +2402,12 @@ async def compile_long_doc( The summary page is already written by the indexer. This function generates concept pages and updates the index. """ - from openkb.config import resolve_effective_config + from openkb.config import resolve_effective_config, resolve_insert_mode config = resolve_effective_config(kb_dir)[0] language: str = config.get("language", "en") entity_types = resolve_entity_types(config) + insert_mode = resolve_insert_mode(config) wiki_dir = kb_dir / "wiki" schema_md = get_agents_md(wiki_dir) @@ -2365,6 +2465,7 @@ async def compile_long_doc( doc_type="pageindex", entity_types=entity_types, bundle=bundle, + insert_mode=insert_mode, ) finally: # Close per-loop litellm async clients before asyncio.run tears this diff --git a/openkb/config.py b/openkb/config.py index 95ca8691f..30533fc1f 100644 --- a/openkb/config.py +++ b/openkb/config.py @@ -36,8 +36,15 @@ # global/KB list overrides it wholesale; resolve_entity_types cleans the # effective value on read. "entity_types": list(DEFAULT_ENTITY_TYPES), + # How a partial compile (concept/entity generation failure) is reported. + # "normal" (today's behavior): warned, file still "added". Strict modes + # raise so the mutation rolls back and the file is "failed" — see + # resolve_insert_mode(). + "insert_mode": "normal", } +VALID_INSERT_MODES: tuple[str, ...] = ("normal", "fail-fast", "fail-at-end") + GLOBAL_CONFIG_DIR = Path.home() / ".config" / "openkb" GLOBAL_CONFIG_PATH = GLOBAL_CONFIG_DIR / "global.yaml" GLOBAL_CONFIG_LOCK_PATH = GLOBAL_CONFIG_DIR / "global.lock" @@ -262,6 +269,25 @@ def resolve_concurrency(config: dict) -> int | None: return value +def resolve_insert_mode(config: dict) -> str: + """Resolve ``insert_mode:`` — one of ``"normal"`` (default, unchanged), + ``"fail-fast"`` (abort on the first concept/entity failure), or + ``"fail-at-end"`` (run to completion, then fail if anything failed). + Strict modes raise ``ConceptCompilationError`` (``openkb.agent.compiler``), + which the mutation rollback turns into a ``"failed"`` outcome. An invalid + value degrades to ``"normal"`` with a warning. + """ + value = config.get("insert_mode", "normal") + if value not in VALID_INSERT_MODES: + logger.warning( + "config: 'insert_mode' must be one of %s, got %r — using 'normal'.", + VALID_INSERT_MODES, + value, + ) + return "normal" + return value + + def resolve_litellm_settings(config: dict) -> dict[str, Any]: """Resolve the optional ``litellm:`` mapping of LiteLLM module settings. diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 95a57cc4c..2404e0d67 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -2085,6 +2085,201 @@ async def test_fallback_list_format(self, tmp_path): assert "Attention" in att_text +class TestInsertMode: + """insert_mode="fail-fast"/"fail-at-end" turn a partial compile (one or + more failed concept/entity generations) into a raised + ConceptCompilationError instead of a silently-partial "success"; see + ``_compile_concepts``'s ``insert_mode`` docstring.""" + + def _setup_wiki(self, tmp_path): + wiki = tmp_path / "wiki" + (wiki / "summaries").mkdir(parents=True) + (wiki / "concepts").mkdir(parents=True) + (wiki / "index.md").write_text( + "# Index\n\n## Documents\n\n## Concepts\n", + encoding="utf-8", + ) + (tmp_path / "raw").mkdir(exist_ok=True) + (tmp_path / "raw" / "test-doc.pdf").write_bytes(b"fake") + return wiki + + @staticmethod + def _selective_acompletion(*args, **kwargs): + """Succeed for "concept-a", raise for "concept-b" — keyed off the + concept title embedded in the page-generation prompt (see + ``_CONCEPT_PAGE_USER``) rather than call order, since concurrent + tasks don't guarantee a fixed completion order.""" + messages = kwargs.get("messages") or (args[1] if len(args) > 1 else []) + # Only the page-specific user message (appended last by _gen_create/ + # _gen_update) carries the concept title — the earlier messages + # (system/doc/summary/known-targets) are shared cached context common + # to every concept in this batch and mention both concept names. + last_content = messages[-1]["content"] if messages else "" + if "concept-b" in last_content: + raise RuntimeError("boom: concept-b generation failed") + mock_resp = MagicMock() + mock_resp.choices = [MagicMock()] + mock_resp.choices[0].message.content = json.dumps( + {"brief": "b", "content": "# Concept A\n\nBody."} + ) + mock_resp.usage = MagicMock(prompt_tokens=100, completion_tokens=50) + mock_resp.usage.prompt_tokens_details = None + return mock_resp + + def _plan_response(self): + return json.dumps( + { + "create": [ + {"name": "concept-a", "title": "concept-a"}, + {"name": "concept-b", "title": "concept-b"}, + ], + "update": [], + "related": [], + } + ) + + @pytest.mark.asyncio + async def test_normal_mode_keeps_partial_success_silent(self, tmp_path): + """Default/omitted insert_mode: unchanged behavior — no exception, + the succeeded concept is written, the failed one is just logged.""" + wiki = self._setup_wiki(tmp_path) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock( + side_effect=_mock_completion([self._plan_response()]) + ) + mock_litellm.acompletion = AsyncMock(side_effect=self._selective_acompletion) + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + insert_mode="normal", + ) + assert (wiki / "concepts" / "concept-a.md").exists() + assert not (wiki / "concepts" / "concept-b.md").exists() + + @pytest.mark.asyncio + async def test_fail_fast_raises_and_writes_nothing(self, tmp_path): + """insert_mode="fail-fast": raises ConceptCompilationError and skips + the write phase entirely — even the concept that already succeeded + must not be written, since the caller rolls back this whole add.""" + from openkb.agent.compiler import ConceptCompilationError + + wiki = self._setup_wiki(tmp_path) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock( + side_effect=_mock_completion([self._plan_response()]) + ) + mock_litellm.acompletion = AsyncMock(side_effect=self._selective_acompletion) + with pytest.raises(ConceptCompilationError): + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + insert_mode="fail-fast", + ) + assert not (wiki / "concepts" / "concept-a.md").exists() + assert not (wiki / "concepts" / "concept-b.md").exists() + + @pytest.mark.asyncio + async def test_fail_at_end_writes_then_raises(self, tmp_path): + """insert_mode="fail-at-end": every planned concept is attempted (so + the succeeded one IS written, unlike fail-fast) before + ConceptCompilationError is raised at the end — the caller's mutation + rollback is what discards the write, not _compile_concepts itself.""" + from openkb.agent.compiler import ConceptCompilationError + + wiki = self._setup_wiki(tmp_path) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock( + side_effect=_mock_completion([self._plan_response()]) + ) + mock_litellm.acompletion = AsyncMock(side_effect=self._selective_acompletion) + with pytest.raises(ConceptCompilationError): + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + insert_mode="fail-at-end", + ) + assert (wiki / "concepts" / "concept-a.md").exists() + assert not (wiki / "concepts" / "concept-b.md").exists() + + @pytest.mark.asyncio + async def test_fail_at_end_no_failures_does_not_raise(self, tmp_path): + """A fully successful compile under insert_mode="fail-at-end" behaves + exactly like "normal" — the strict check only fires on an actual + failure.""" + wiki = self._setup_wiki(tmp_path) + plan_response = json.dumps( + {"create": [{"name": "concept-a", "title": "concept-a"}], "update": [], "related": []} + ) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock(side_effect=_mock_completion([plan_response])) + mock_litellm.acompletion = AsyncMock(side_effect=self._selective_acompletion) + await _compile_concepts( + wiki, + tmp_path, + "gpt-4o-mini", + {"role": "system", "content": "s"}, + {"role": "user", "content": "d"}, + "summary", + "test-doc", + 5, + insert_mode="fail-at-end", + ) + assert (wiki / "concepts" / "concept-a.md").exists() + + @pytest.mark.asyncio + async def test_compile_short_doc_reads_insert_mode_from_kb_config(self, tmp_path): + """End-to-end wiring check: compile_short_doc reads insert_mode from + the KB's config.yaml (via resolve_effective_config) and threads it + into _compile_concepts — no CLI-level plumbing needed.""" + from openkb.agent.compiler import ConceptCompilationError + + wiki = tmp_path / "wiki" + (wiki / "sources").mkdir(parents=True) + (wiki / "summaries").mkdir(parents=True) + (wiki / "concepts").mkdir(parents=True) + (wiki / "index.md").write_text( + "# Index\n\n## Documents\n\n## Concepts\n", + encoding="utf-8", + ) + source_path = wiki / "sources" / "test-doc.md" + source_path.write_text("# Test Doc\n\nContent.", encoding="utf-8") + openkb_dir = tmp_path / ".openkb" + openkb_dir.mkdir() + (openkb_dir / "config.yaml").write_text("insert_mode: fail-fast\n", encoding="utf-8") + (tmp_path / "raw").mkdir() + (tmp_path / "raw" / "test-doc.pdf").write_bytes(b"fake") + + summary_response = json.dumps({"description": "d", "content": "# Summary\n\nContent."}) + with patch("openkb.agent.compiler.litellm") as mock_litellm: + mock_litellm.completion = MagicMock( + side_effect=_mock_completion([summary_response, self._plan_response()]) + ) + mock_litellm.acompletion = AsyncMock(side_effect=self._selective_acompletion) + with pytest.raises(ConceptCompilationError): + await compile_short_doc("test-doc", source_path, tmp_path, "gpt-4o-mini") + # fail-fast: neither concept must have been written. + assert not (wiki / "concepts" / "concept-a.md").exists() + assert not (wiki / "concepts" / "concept-b.md").exists() + + class TestBriefIntegration: @pytest.mark.asyncio async def test_short_doc_briefs_in_index_and_frontmatter(self, tmp_path): diff --git a/tests/test_config.py b/tests/test_config.py index 65572d6b9..c4f4f3905 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -6,6 +6,7 @@ from openkb.config import ( DEFAULT_CONFIG, GLOBAL_SCALAR_KEYS, + VALID_INSERT_MODES, get_extra_headers, get_parallel_tool_calls, get_timeout, @@ -16,6 +17,7 @@ resolve_effective_config, resolve_extra_headers, resolve_init_kb_dir, + resolve_insert_mode, resolve_litellm_settings, resolve_model_settings, resolve_parallel_tool_calls, @@ -196,6 +198,23 @@ def test_resolve_concurrency_none_is_silent(caplog): assert caplog.text == "" +def test_resolve_insert_mode_absent_is_normal(): + assert resolve_insert_mode({}) == "normal" + + +def test_resolve_insert_mode_valid_values(): + assert set(VALID_INSERT_MODES) == {"normal", "fail-fast", "fail-at-end"} + for mode in VALID_INSERT_MODES: + assert resolve_insert_mode({"insert_mode": mode}) == mode + + +def test_resolve_insert_mode_rejects_invalid(caplog): + with caplog.at_level(logging.WARNING, logger="openkb.config"): + result = resolve_insert_mode({"insert_mode": "yolo"}) + assert result == "normal" + assert "insert_mode" in caplog.text + + def test_load_missing_file_returns_defaults(tmp_path): missing = tmp_path / "nonexistent" / "config.yaml" config = load_config(missing) From 153149cff0b632391f935c6930bb11401ca70b54 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Mon, 31 Aug 2026 10:43:44 +0200 Subject: [PATCH 16/16] feat(cli): add empty subdirectory cleanup to add-all command When auto_delete_added_files is enabled, the add-all command now recursively cleans up empty subdirectories in raw/ after file deletion. Directories are deleted from deepest to shallowest to ensure proper cleanup. Added summary output showing count of cleaned directories. --- openkb/cli.py | 43 ++++++++++++++++++++++++++++++++++++++----- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/openkb/cli.py b/openkb/cli.py index 8e6086523..7a457f1cc 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -549,6 +549,33 @@ def _delete_if_auto_cleanup_enabled( return False +def _cleanup_empty_directories(start_dir: Path) -> int: + """Recursively delete empty directories under start_dir. + + Walks from deepest subdirectories up, deleting directories that become + empty after file cleanup. + + Args: + start_dir: Root directory to clean up (e.g., kb_dir / "raw"). + + Returns: + Number of directories deleted. + """ + deleted_count = 0 + try: + for directory in sorted(start_dir.rglob("*"), key=lambda p: len(p.parts), reverse=True): + if directory.is_dir() and directory != start_dir: + try: + if not list(directory.iterdir()): + directory.rmdir() + deleted_count += 1 + except OSError: + pass + except Exception as exc: + logger.warning(f"Error during directory cleanup: {exc}") + return deleted_count + + def _add_single_file_locked( file_path: Path, kb_dir: Path, *, stage: bool = True, bundle=None ) -> Literal["added", "skipped", "failed"]: @@ -1301,7 +1328,8 @@ def add_all(ctx): This command walks the ``raw/`` directory recursively for all supported document types and ingests them into the KB. If ``auto_delete_added_files`` is enabled in config.yaml, files are automatically deleted after ingestion - (both on successful addition and on skip/duplicate). + (both on successful addition and on skip/duplicate), and empty subdirectories + are cleaned up. Returns a summary of the operation (added, skipped, failed, deleted counts). """ @@ -1326,7 +1354,7 @@ def add_all(ctx): config = resolve_effective_config(kb_dir)[0] total = len(files) - added = skipped = failed = deleted = 0 + added = skipped = failed = deleted = dirs_deleted = 0 click.echo(f"Processing {total} file(s) from raw/ directory...") for i, f in enumerate(files, 1): @@ -1341,9 +1369,14 @@ def add_all(ctx): if _delete_if_auto_cleanup_enabled(f, outcome, config): deleted += 1 - click.echo( - f"\n\nSummary: Added: {added}, Skipped: {skipped}, Failed: {failed}, Deleted: {deleted}" - ) + # Clean up empty subdirectories if auto-cleanup is enabled + if config.get("auto_delete_added_files", False): + dirs_deleted = _cleanup_empty_directories(raw_dir) + + summary = f"Added: {added}, Skipped: {skipped}, Failed: {failed}, Deleted: {deleted}" + if dirs_deleted > 0: + summary += f", Empty dirs cleaned: {dirs_deleted}" + click.echo(f"\n\nSummary: {summary}") def _stream_to_tty() -> bool: