From 9180f4d91045ebad8e1696dd81aebccad7972c26 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Thu, 27 Aug 2026 15:27:26 +0200 Subject: [PATCH 1/3] 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 2/3] 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 3/3] 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). """