refactor: dedupe LLM resilience/ingestion retries, type RAG source boundary

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":

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)."""

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,14 +48,20 @@ 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,
             status=JobStatus.PENDING,
             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

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,

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))

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]:

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)

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))