From a4ba6ead7325c7d50619ef82d084a491bc1b5b29 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 21:11:30 +0200 Subject: [PATCH] refactor: dedupe LLM resilience/ingestion retry loops, type RAG source boundary - kill 4 type: ignore casts on the RAG->API source boundary via SourceRef(**s) - collapse ResilientProvider's 4 duplicated breaker/fallback blocks into _call_with_breaker/_stream_with_breaker - collapse ingestion's two-pass retry loop into a shared _ingest_batch helper - type IngestionService.bm25_index as Bm25Index | None instead of Any | None - add AgentApiError so /agent/query follows the same error-mapping convention as every other router - 404 DELETE /collections/{name} on unknown name instead of silently no-oping - cap concurrent ingest job submission (max_pending_ingest_jobs, default 10), reject past-cap submissions with 429 instead of unbounded queuing Claude-Session: https://claude.ai/code/session_01NFJtv7kyDmx4W8u76WW9EC --- localrag/agent/service.py | 4 +- localrag/api/exceptions.py | 4 ++ localrag/api/jobs.py | 12 ++++- localrag/api/routers/agent.py | 10 ++--- localrag/api/service.py | 20 ++++----- localrag/ingestion/service.py | 80 +++++++++++++++++++-------------- localrag/llm/resilience.py | 85 +++++++++++++++++++---------------- localrag/rag/engine.py | 4 +- localrag/settings.py | 3 ++ tests/test_api_jobs.py | 18 +++++++- tests/test_api_service.py | 4 ++ 11 files changed, 149 insertions(+), 95 deletions(-) diff --git a/localrag/agent/service.py b/localrag/agent/service.py index 9d51103..e167440 100644 --- a/localrag/agent/service.py +++ b/localrag/agent/service.py @@ -76,7 +76,7 @@ class AgentResponse: answer: str tool_used: str reasoning: str - sources: list[dict[str, object]] = field(default_factory=list) + sources: list[dict[str, Any]] = field(default_factory=list) latency_ms: float = 0.0 model: str = "" @@ -102,7 +102,7 @@ def run_agent( tool_used = "unknown" reasoning = "" answer = "" - sources: list[dict[str, object]] = [] + sources: list[dict[str, Any]] = [] for block in resp.content: if block.type != "tool_use": diff --git a/localrag/api/exceptions.py b/localrag/api/exceptions.py index 1072db2..9410c2e 100644 --- a/localrag/api/exceptions.py +++ b/localrag/api/exceptions.py @@ -16,3 +16,7 @@ class IngestApiError(HttpMappedError): class RagApiError(HttpMappedError): """Raised when RAG query cannot complete (embedding or vector store failure).""" + + +class AgentApiError(HttpMappedError): + """Raised when the agent endpoint cannot run (e.g. missing provider credentials).""" diff --git a/localrag/api/jobs.py b/localrag/api/jobs.py index d752639..2c1299a 100644 --- a/localrag/api/jobs.py +++ b/localrag/api/jobs.py @@ -27,6 +27,10 @@ class JobStatus(StrEnum): FAILED = "failed" +class TooManyPendingJobsError(Exception): + """Raised when a submission would exceed the configured pending/running job cap.""" + + @dataclass class Job: job_id: str @@ -44,7 +48,7 @@ class JobRegistry: default_factory=lambda: ThreadPoolExecutor(max_workers=2, thread_name_prefix="ingest-job") ) - def submit(self, work: Callable[[], dict[str, Any]]) -> str: + def submit(self, work: Callable[[], dict[str, Any]], *, max_pending: int = 10) -> str: job_id = uuid4().hex job = Job( job_id=job_id, @@ -52,6 +56,12 @@ def submit(self, work: Callable[[], dict[str, Any]]) -> str: created_at=datetime.now(UTC).isoformat(), ) with self._lock: + pending = sum( + 1 for j in self._jobs.values() if j.status in (JobStatus.PENDING, JobStatus.RUNNING) + ) + if pending >= max_pending: + message = f"{pending} ingest jobs already pending/running (max {max_pending})." + raise TooManyPendingJobsError(message) self._jobs[job_id] = job self._executor.submit(self._run, job_id, work) return job_id diff --git a/localrag/api/routers/agent.py b/localrag/api/routers/agent.py index adf6733..2c8bbe3 100644 --- a/localrag/api/routers/agent.py +++ b/localrag/api/routers/agent.py @@ -2,10 +2,11 @@ from http import HTTPStatus -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends from localrag.agent.service import run_agent from localrag.api.dependencies import get_api_settings, get_engine, require_api_key +from localrag.api.exceptions import AgentApiError from localrag.api.schemas import AgentQueryRequest, AgentQueryResponse, SourceRef from localrag.rag.engine import RAGEngine from localrag.settings import Settings @@ -24,7 +25,7 @@ def agent_query( The agent decides whether to search documents or answer directly. """ if not settings.anthropic_api_key: - raise HTTPException( + raise AgentApiError( status_code=HTTPStatus.SERVICE_UNAVAILABLE, detail="ANTHROPIC_API_KEY is not configured. Set it in .env to use the agent endpoint.", ) @@ -35,10 +36,7 @@ def agent_query( api_key=settings.anthropic_api_key, model=model, ) - sources = [ - SourceRef(source=str(s.get("source", "")), chunk_index=int(s.get("chunk_index", -1))) # type: ignore[call-overload] - for s in result.sources - ] + sources = [SourceRef(**s) for s in result.sources] return AgentQueryResponse( answer=result.answer, tool_used=result.tool_used, diff --git a/localrag/api/service.py b/localrag/api/service.py index 29caf3e..11de87c 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -14,7 +14,7 @@ from localrag import metrics as app_metrics from localrag.api.exceptions import IngestApiError, RagApiError -from localrag.api.jobs import JobRegistry +from localrag.api.jobs import JobRegistry, TooManyPendingJobsError from localrag.api.repository import ChromaCollectionRepository from localrag.api.schemas import ( CollectionDeleteResponse, @@ -89,6 +89,8 @@ def list_collections_response( def delete_collection_response( collection_repo: ChromaCollectionRepository, name: str ) -> CollectionDeleteResponse: + if name not in collection_repo.list_collection_names(): + raise IngestApiError(HTTPStatus.NOT_FOUND, f"Collection '{name}' not found.") logger.warning("collection_delete name=%s", name) collection_repo.delete_collection(name) return CollectionDeleteResponse(status="ok") @@ -234,7 +236,10 @@ def work() -> dict[str, Any]: ], } - job_id = job_registry.submit(work) + try: + job_id = job_registry.submit(work, max_pending=settings.max_pending_ingest_jobs) + except TooManyPendingJobsError as exc: + raise IngestApiError(HTTPStatus.TOO_MANY_REQUESTS, str(exc)) from exc logger.info("ingest_directory_async_submitted job_id=%s path=%s", job_id, path) return IngestJobResponse(job_id=job_id, status="pending") @@ -356,16 +361,7 @@ def query_json( latency_ms = (time.perf_counter() - t0) * 1000 used_model = request.model or engine.settings.ollama_llm_model - raw_sources = engine.extract_sources(contexts) - sources = [ - SourceRef( - source=str(s.get("source", "")), - chunk_index=int(s.get("chunk_index", -1)), # type: ignore[call-overload] - heading_path=s.get("heading_path"), # type: ignore[arg-type] - chunk_type=s.get("chunk_type"), # type: ignore[arg-type] - ) - for s in raw_sources - ] + sources = [SourceRef(**s) for s in engine.extract_sources(contexts)] app_metrics.query_duration_seconds.observe(latency_ms / 1000) app_metrics.chunks_retrieved_total.inc(len(contexts)) diff --git a/localrag/ingestion/service.py b/localrag/ingestion/service.py index a8bfd29..3175eea 100644 --- a/localrag/ingestion/service.py +++ b/localrag/ingestion/service.py @@ -3,15 +3,16 @@ import hashlib import logging import subprocess +from collections.abc import Callable from dataclasses import dataclass, field from datetime import UTC, datetime from pathlib import Path -from typing import Any from localrag.ingestion.chunker import chunk_text from localrag.ingestion.embedder import OllamaEmbedder from localrag.ingestion.loader import list_supported_files, parse_file from localrag.ingestion.structural_chunker import Chunk, chunk_document +from localrag.rag.bm25_index import Bm25Index from localrag.settings import Settings, is_path_allowed from localrag.storage.vector_store import VectorStore @@ -51,7 +52,7 @@ class IngestionService: settings: Settings embedder: OllamaEmbedder vector_store: VectorStore - bm25_index: Any | None = None + bm25_index: Bm25Index | None = None def ingest_file(self, path: Path, embed_model: str | None = None) -> IngestionResult: return self.ingest_paths([path], embed_model=embed_model) @@ -108,41 +109,29 @@ def _stored_content_hashes(self) -> dict[str, str]: return hashes def ingest_paths(self, paths: list[Path], embed_model: str | None = None) -> IngestionResult: - files_processed = 0 - total_chunks = 0 - processed_sources: list[str] = [] - retry_queue: list[Path] = [] - - for resolved_path in self._allowed_paths(paths): - chunks_added, error = self._safe_ingest_one(resolved_path, embed_model) - if error is not None: - logger.warning( - "ingest_file_failed_will_retry path=%s error=%s", resolved_path, error - ) - retry_queue.append(resolved_path) - continue - if chunks_added is None: - continue - files_processed += 1 - total_chunks += chunks_added - processed_sources.append(str(resolved_path)) + files_processed, total_chunks, processed_sources, retry_queue = self._ingest_batch( + self._allowed_paths(paths), + embed_model, + log_fn=logger.warning, + log_event="ingest_file_failed_will_retry", + ) failed_sources: list[FailedSource] = [] if retry_queue: logger.info("ingest_retry_start count=%s", len(retry_queue)) - for resolved_path in retry_queue: - chunks_added, error = self._safe_ingest_one(resolved_path, embed_model) - if error is not None: - logger.error( - "ingest_file_failed_permanently path=%s error=%s", resolved_path, error - ) - failed_sources.append(FailedSource(source=str(resolved_path), error=error)) - continue - if chunks_added is None: - continue - files_processed += 1 - total_chunks += chunks_added - processed_sources.append(str(resolved_path)) + retry_paths = [path for path, _error in retry_queue] + retry_processed, retry_chunks, retry_sources, still_failed = self._ingest_batch( + retry_paths, + embed_model, + log_fn=logger.error, + log_event="ingest_file_failed_permanently", + ) + files_processed += retry_processed + total_chunks += retry_chunks + processed_sources.extend(retry_sources) + failed_sources.extend( + FailedSource(source=str(path), error=error) for path, error in still_failed + ) # Rebuilding the BM25 corpus is O(total chunks); do it once per batch, not per file. if self.bm25_index is not None and files_processed > 0: @@ -155,6 +144,31 @@ def ingest_paths(self, paths: list[Path], embed_model: str | None = None) -> Ing failed_sources=failed_sources, ) + def _ingest_batch( + self, + paths: list[Path], + embed_model: str | None, + *, + log_fn: Callable[..., None], + log_event: str, + ) -> tuple[int, int, list[str], list[tuple[Path, str]]]: + files_processed = 0 + total_chunks = 0 + processed_sources: list[str] = [] + failed: list[tuple[Path, str]] = [] + for resolved_path in paths: + chunks_added, error = self._safe_ingest_one(resolved_path, embed_model) + if error is not None: + log_fn("%s path=%s error=%s", log_event, resolved_path, error) + failed.append((resolved_path, error)) + continue + if chunks_added is None: + continue + files_processed += 1 + total_chunks += chunks_added + processed_sources.append(str(resolved_path)) + return files_processed, total_chunks, processed_sources, failed + def _safe_ingest_one( self, resolved_path: Path, embed_model: str | None ) -> tuple[int | None, str | None]: diff --git a/localrag/llm/resilience.py b/localrag/llm/resilience.py index 8c13efe..9fdf3a6 100644 --- a/localrag/llm/resilience.py +++ b/localrag/llm/resilience.py @@ -10,6 +10,7 @@ import logging from collections.abc import Callable, Generator +from functools import partial from typing import Any, TypeVar import anthropic @@ -70,50 +71,63 @@ def _retrying(self, func: Callable[[], _T]) -> _T: )(func) return wrapped() - def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + def _call_with_breaker(self, fn: Callable[[], _T], fallback: Callable[[], _T] | None) -> _T: try: - return self._breaker.call( - self._retrying, lambda: self._provider.generate(prompt, context, model=model) - ) + return self._breaker.call(self._retrying, fn) except pybreaker.CircuitBreakerError: - logger.warning("llm_circuit_open falling_back=%s", self._fallback_provider is not None) - if self._fallback_provider is not None: - return self._fallback_provider.generate(prompt, context, model=model) + logger.warning("llm_circuit_open falling_back=%s", fallback is not None) + if fallback is not None: + return fallback() raise - def stream( - self, prompt: str, context: list[str], *, model: str | None = None + def _stream_with_breaker( + self, + start_fn: Callable[[], tuple[Generator[dict[str, Any]], dict[str, Any] | None]], + fallback_stream: Callable[[], Generator[dict[str, Any]]] | None, ) -> Generator[dict[str, Any]]: # Streaming can't be retried mid-token-stream: the retry/breaker call only # guards establishing the stream (forcing the first event to surface any # connection error), then remaining tokens pass through untouched. - def start() -> tuple[Generator[dict[str, Any]], dict[str, Any] | None]: - gen = self._provider.stream(prompt, context, model=model) - first_event = next(gen, None) - return gen, first_event - try: - gen, first_event = self._breaker.call(self._retrying, start) + gen, first_event = self._breaker.call(self._retrying, start_fn) except pybreaker.CircuitBreakerError: - logger.warning("llm_circuit_open falling_back=%s", self._fallback_provider is not None) - if self._fallback_provider is not None: - yield from self._fallback_provider.stream(prompt, context, model=model) + logger.warning("llm_circuit_open falling_back=%s", fallback_stream is not None) + if fallback_stream is not None: + yield from fallback_stream() return raise if first_event is not None: yield first_event yield from gen + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + fallback = None + if self._fallback_provider is not None: + fallback = partial(self._fallback_provider.generate, prompt, context, model=model) + return self._call_with_breaker( + partial(self._provider.generate, prompt, context, model=model), fallback + ) + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + def start() -> tuple[Generator[dict[str, Any]], dict[str, Any] | None]: + gen = self._provider.stream(prompt, context, model=model) + first_event = next(gen, None) + return gen, first_event + + fallback_stream = None + if self._fallback_provider is not None: + fallback_stream = partial(self._fallback_provider.stream, prompt, context, model=model) + yield from self._stream_with_breaker(start, fallback_stream) + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: - try: - return self._breaker.call( - self._retrying, lambda: self._provider.generate_from_prompt(prompt, model=model) - ) - except pybreaker.CircuitBreakerError: - logger.warning("llm_circuit_open falling_back=%s", self._fallback_provider is not None) - if self._fallback_provider is not None: - return self._fallback_provider.generate_from_prompt(prompt, model=model) - raise + fallback = None + if self._fallback_provider is not None: + fallback = partial(self._fallback_provider.generate_from_prompt, prompt, model=model) + return self._call_with_breaker( + partial(self._provider.generate_from_prompt, prompt, model=model), fallback + ) def stream_from_prompt( self, prompt: str, *, model: str | None = None @@ -123,17 +137,12 @@ def start() -> tuple[Generator[dict[str, Any]], dict[str, Any] | None]: first_event = next(gen, None) return gen, first_event - try: - gen, first_event = self._breaker.call(self._retrying, start) - except pybreaker.CircuitBreakerError: - logger.warning("llm_circuit_open falling_back=%s", self._fallback_provider is not None) - if self._fallback_provider is not None: - yield from self._fallback_provider.stream_from_prompt(prompt, model=model) - return - raise - if first_event is not None: - yield first_event - yield from gen + fallback_stream = None + if self._fallback_provider is not None: + fallback_stream = partial( + self._fallback_provider.stream_from_prompt, prompt, model=model + ) + yield from self._stream_with_breaker(start, fallback_stream) def count_tokens(self, text: str) -> int: return self._provider.count_tokens(text) diff --git a/localrag/rag/engine.py b/localrag/rag/engine.py index 5d516c4..3ac9d10 100644 --- a/localrag/rag/engine.py +++ b/localrag/rag/engine.py @@ -105,9 +105,9 @@ def _stream_chat_tokens( yield {"type": "final", "sources": self.extract_sources(contexts), "low_confidence": False} @staticmethod - def extract_sources(contexts: list[dict[str, Any]]) -> list[dict[str, object]]: + def extract_sources(contexts: list[dict[str, Any]]) -> list[dict[str, Any]]: seen: set[tuple[str, int]] = set() - sources: list[dict[str, object]] = [] + sources: list[dict[str, Any]] = [] for context in contexts: source = str(context.get("source", "unknown")) chunk_index = int(context.get("chunk_index", -1)) diff --git a/localrag/settings.py b/localrag/settings.py index e99c3a1..33f1da1 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -98,6 +98,9 @@ class Settings(BaseSettings): ingest_recursive: bool = True ingest_roots: list[str] = [] + # Caps concurrently pending/running async ingest jobs; further submissions get 429. + max_pending_ingest_jobs: int = 10 + upload_dir: str = "./data/uploads" upload_max_bytes: int = 100_000_000 diff --git a/tests/test_api_jobs.py b/tests/test_api_jobs.py index 3a7c000..eec02b5 100644 --- a/tests/test_api_jobs.py +++ b/tests/test_api_jobs.py @@ -1,8 +1,11 @@ from __future__ import annotations +import threading import time -from localrag.api.jobs import Job, JobRegistry, JobStatus +import pytest + +from localrag.api.jobs import Job, JobRegistry, JobStatus, TooManyPendingJobsError def _wait_for_terminal(registry: JobRegistry, job_id: str, timeout: float = 5.0) -> Job | None: @@ -44,3 +47,16 @@ def boom() -> dict[str, object]: def test_job_registry_get_returns_none_for_unknown_job() -> None: registry = JobRegistry() assert registry.get("nope") is None + + +def test_job_registry_submit_rejects_past_pending_cap() -> None: + registry = JobRegistry() + release = threading.Event() + job_ids = [registry.submit(release.wait, max_pending=2) for _ in range(2)] + + with pytest.raises(TooManyPendingJobsError): + registry.submit(lambda: {"value": 1}, max_pending=2) + + release.set() + for job_id in job_ids: + _wait_for_terminal(registry, job_id) diff --git a/tests/test_api_service.py b/tests/test_api_service.py index 1dd5369..6294236 100644 --- a/tests/test_api_service.py +++ b/tests/test_api_service.py @@ -69,6 +69,10 @@ def test_health_success_and_collections_endpoints(tmp_path: Path) -> None: assert deleted.json() == {"status": "ok"} assert repo.deleted == ["col-1"] + missing = client.delete("/collections/does-not-exist") + assert missing.status_code == 404 + assert repo.deleted == ["col-1"] + app.dependency_overrides.clear()