From 2f91eff5ceaf3514a0c576aab88e4b6ae7b9d0f1 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 18:23:02 +0200 Subject: [PATCH 01/20] feat: content-hash-based incremental rebuild + richer ingestion metadata --- docs/rag-retrieval.md | 6 +++ localrag/api/schemas.py | 8 ++++ localrag/api/service.py | 1 + localrag/ingestion/service.py | 57 ++++++++++++++++++++++---- tests/test_ingestion_service.py | 72 ++++++++++++++++++++++++++++++++- 5 files changed, 136 insertions(+), 8 deletions(-) diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index 88b949a..84eb058 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -54,6 +54,12 @@ Freshness and debugging depend on chunk metadata written during ingestion: - `ingested_at` - `heading_path` - `chunk_type` +- `content_hash` +- `source_mtime` +- `git_commit` The retriever returns `freshness_factor` and `ingested_at` in contexts so rank decisions are visible in API and test traces. + +`content_hash` also drives incremental rebuild — `POST /collections/rebuild` skips +re-embedding any source whose file bytes haven't changed. diff --git a/localrag/api/schemas.py b/localrag/api/schemas.py index 063ece3..e490aaf 100644 --- a/localrag/api/schemas.py +++ b/localrag/api/schemas.py @@ -343,3 +343,11 @@ class RebuildCollectionResponse(BaseModel): "Their old vectors remain in place (not removed); also logged at ERROR level." ), ) + skipped_unchanged_sources: list[str] = Field( + default_factory=list, + description=( + "Sources whose file content hash matched the previously stored chunk metadata; " + "skipped entirely (no re-parse/re-embed) by the incremental rebuild." + ), + examples=[["/docs/unchanged.md"]], + ) diff --git a/localrag/api/service.py b/localrag/api/service.py index fdeadab..de1a737 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -109,6 +109,7 @@ def rebuild_collection_response( failed_sources=[ FailedSourceRef(source=f.source, error=f.error) for f in result.failed_sources ], + skipped_unchanged_sources=result.skipped_unchanged_sources, ) diff --git a/localrag/ingestion/service.py b/localrag/ingestion/service.py index e1c9f3a..d43a5dc 100644 --- a/localrag/ingestion/service.py +++ b/localrag/ingestion/service.py @@ -1,6 +1,8 @@ from __future__ import annotations +import hashlib import logging +import subprocess from dataclasses import dataclass, field from datetime import UTC, datetime from pathlib import Path @@ -41,6 +43,7 @@ class RebuildCollectionResult: processed_sources: list[str] missing_sources: list[str] failed_sources: list[FailedSource] = field(default_factory=list) + skipped_unchanged_sources: list[str] = field(default_factory=list) @dataclass @@ -68,28 +71,42 @@ def ingest_directory( def rebuild_collection(self, embed_model: str | None = None) -> RebuildCollectionResult: sources = self.vector_store.list_distinct_sources() + stored_hashes = self._stored_content_hashes() missing_sources: list[str] = [] + skipped_unchanged: list[str] = [] paths_to_ingest: list[Path] = [] for source in sources: path = Path(source) - if path.is_file(): - paths_to_ingest.append(path) - else: + if not path.is_file(): missing_sources.append(source) self.vector_store.delete_by_source(source) - logger.warning( - "rebuild_skip_missing_file source=%s", - source, - ) + logger.warning("rebuild_skip_missing_file source=%s", source) + continue + current_hash = _file_content_hash(path) + if stored_hashes.get(source) == current_hash: + skipped_unchanged.append(source) + continue + paths_to_ingest.append(path) ingest = self.ingest_paths(paths_to_ingest, embed_model=embed_model) + logger.info("rebuild_skipped_unchanged count=%s", len(skipped_unchanged)) return RebuildCollectionResult( files_processed=ingest.files_processed, total_chunks=ingest.total_chunks, processed_sources=ingest.processed_sources, missing_sources=sorted(missing_sources), failed_sources=ingest.failed_sources, + skipped_unchanged_sources=sorted(skipped_unchanged), ) + def _stored_content_hashes(self) -> dict[str, str]: + hashes: dict[str, str] = {} + for _chunk_id, _document, metadata in self.vector_store.get_all_chunks(): + source = metadata.get("source") + content_hash = metadata.get("content_hash") + if isinstance(source, str) and isinstance(content_hash, str) and source not in hashes: + hashes[source] = content_hash + return hashes + def ingest_paths(self, paths: list[Path], embed_model: str | None = None) -> IngestionResult: total_chunks = 0 files_processed = 0 @@ -172,6 +189,9 @@ def _ingest_one(self, resolved_path: Path, embed_model: str | None) -> int | Non # embed call leaves the previous (still valid) vectors for this source in place. self.vector_store.delete_by_source(source) created_at = datetime.now(UTC).isoformat() + content_hash = _file_content_hash(resolved_path) + source_mtime = resolved_path.stat().st_mtime + git_commit = _git_commit_for_path(resolved_path) or "" metadatas = [ { "source": source, @@ -180,6 +200,9 @@ def _ingest_one(self, resolved_path: Path, embed_model: str | None) -> int | Non "heading_path": chunk.heading_path, "chunk_type": chunk.chunk_type, "ingested_at": created_at, + "content_hash": content_hash, + "source_mtime": source_mtime, + "git_commit": git_commit, } for index, chunk in enumerate(structural_chunks) ] @@ -203,3 +226,23 @@ def _build_chunks(self, text: str, file_type: str) -> list[Chunk]: Chunk(text=chunk, heading_path="", chunk_type="fixed") for chunk in fixed_chunks ] return chunk_document(text=text, file_type=file_type, settings=self.settings) + + +def _file_content_hash(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _git_commit_for_path(path: Path) -> str | None: + try: + result = subprocess.run( # noqa: S603 + ["git", "log", "-1", "--format=%H", "--", path.name], # noqa: S607 + cwd=path.parent, + capture_output=True, + text=True, + timeout=5, + check=False, + ) + except (OSError, subprocess.TimeoutExpired): + return None + commit = result.stdout.strip() + return commit or None diff --git a/tests/test_ingestion_service.py b/tests/test_ingestion_service.py index de4ad28..34d9ec3 100644 --- a/tests/test_ingestion_service.py +++ b/tests/test_ingestion_service.py @@ -1,6 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass +import hashlib +from dataclasses import dataclass, field from pathlib import Path import pytest @@ -26,6 +27,7 @@ class StubVectorStore: deleted_sources: list[str] added: list[dict[str, object]] distinct_sources: list[str] | None = None + all_chunks: list[tuple[str, str, dict[str, object]]] = field(default_factory=list) def list_distinct_sources(self) -> list[str]: return list(self.distinct_sources or []) @@ -33,6 +35,9 @@ def list_distinct_sources(self) -> list[str]: def delete_by_source(self, source: str) -> None: self.deleted_sources.append(source) + def get_all_chunks(self) -> list[tuple[str, str, dict[str, object]]]: + return list(self.all_chunks) + def add_chunks( self, source: str, @@ -301,3 +306,68 @@ def test_ingestion_service_rebuild_reingests_distinct_sources(tmp_path: Path) -> assert vector_store.deleted_sources[0] == missing assert str(kept.resolve()) in result.processed_sources assert result.files_processed == 1 + + +def test_ingestion_service_ingest_one_writes_content_hash_and_git_metadata(tmp_path: Path) -> None: + allowed_root = tmp_path / "allowed" + allowed_root.mkdir() + path = allowed_root / "a.md" + path.write_text("hello world", encoding="utf-8") + + settings = Settings(ingest_roots=[str(allowed_root)], chunk_chars=100, chunk_overlap_chars=0) + embedder = StubEmbedder(seen_texts_batches=[]) + vector_store = StubVectorStore(deleted_sources=[], added=[], distinct_sources=None) + service = IngestionService( + settings=settings, + embedder=embedder, # type: ignore[arg-type] + vector_store=vector_store, # type: ignore[arg-type] + ) + + service.ingest_paths([path]) + + added = vector_store.added[0] + metadatas = added["metadatas"] + expected_hash = hashlib.sha256(path.read_bytes()).hexdigest() + assert all(md["content_hash"] == expected_hash for md in metadatas) # type: ignore[union-attr] + assert all("source_mtime" in md for md in metadatas) # type: ignore[union-attr] + assert all(md["git_commit"] == "" for md in metadatas) # type: ignore[union-attr] + + +def test_ingestion_service_rebuild_skips_unchanged_sources(tmp_path: Path) -> None: + allowed_root = tmp_path / "allowed" + allowed_root.mkdir() + unchanged = allowed_root / "unchanged.md" + unchanged.write_text("same content", encoding="utf-8") + changed = allowed_root / "changed.md" + changed.write_text("new content", encoding="utf-8") + + settings = Settings(ingest_roots=[str(allowed_root)], chunk_chars=100, chunk_overlap_chars=0) + embedder = StubEmbedder(seen_texts_batches=[]) + unchanged_hash = hashlib.sha256(unchanged.read_bytes()).hexdigest() + vector_store = StubVectorStore( + deleted_sources=[], + added=[], + distinct_sources=[str(unchanged.resolve()), str(changed.resolve())], + all_chunks=[ + ( + "id-1", + "same content", + {"source": str(unchanged.resolve()), "content_hash": unchanged_hash}, + ), + ( + "id-2", + "old content", + {"source": str(changed.resolve()), "content_hash": "stale-hash"}, + ), + ], + ) + service = IngestionService( + settings=settings, + embedder=embedder, # type: ignore[arg-type] + vector_store=vector_store, # type: ignore[arg-type] + ) + + result = service.rebuild_collection() + + assert result.skipped_unchanged_sources == [str(unchanged.resolve())] + assert result.processed_sources == [str(changed.resolve())] From 6d038504bb53142e02d6037f341debe29f61c62e Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 21:22:49 +0200 Subject: [PATCH 02/20] fix: bump default chunk overlap to ~12.5% of max chunk size --- .env.example | 2 +- docs/rag-retrieval.md | 10 ++++++++++ localrag/settings.py | 2 +- tests/test_settings.py | 9 +++++++++ 4 files changed, 21 insertions(+), 2 deletions(-) create mode 100644 tests/test_settings.py diff --git a/.env.example b/.env.example index 660f96e..b094793 100644 --- a/.env.example +++ b/.env.example @@ -6,7 +6,7 @@ CHROMA_PERSIST_PATH=./data/chroma CHROMA_COLLECTION_NAME=localrag CHUNK_CHARS=512 -CHUNK_OVERLAP_CHARS=50 +CHUNK_OVERLAP_CHARS=150 CHUNKING_MODE=structural CHUNK_MAX_CHARS=1200 CHUNK_MIN_CHARS=200 diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index 84eb058..c58b2f7 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -45,6 +45,16 @@ Final score becomes: Set `FRESHNESS_HALF_LIFE_DAYS=0` to disable this behavior. +## Chunk overlap + +`CHUNK_OVERLAP_CHARS` (default 150, ~12.5% of `CHUNK_MAX_CHARS=1200`) only applies +where `localrag/ingestion/structural_chunker.py::_split_long_paragraph` must +hard-split a single paragraph that exceeds `chunk_max_chars`. Adjacent +*packed* structural chunks (the common case — `_pack_blocks`) are +deliberately disjoint with zero overlap: boundary-awareness (never splitting +mid-table, mid-code-block, or mid-heading-section) substitutes for overlap +there. This is an intentional design choice, not an oversight. + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: diff --git a/localrag/settings.py b/localrag/settings.py index 5f36779..fbeffaa 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -68,7 +68,7 @@ class Settings(BaseSettings): chroma_collection_name: str = "localrag" chunk_chars: int = 512 - chunk_overlap_chars: int = 50 + chunk_overlap_chars: int = 150 chunking_mode: str = "structural" chunk_max_chars: int = 1200 chunk_min_chars: int = 200 diff --git a/tests/test_settings.py b/tests/test_settings.py new file mode 100644 index 0000000..ca4654f --- /dev/null +++ b/tests/test_settings.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +from localrag.settings import Settings + + +def test_default_chunk_overlap_is_within_10_to_20_percent_of_max_chars() -> None: + settings = Settings() + ratio = settings.chunk_overlap_chars / settings.chunk_max_chars + assert 0.10 <= ratio <= 0.20 From d7c1ef5ff9b17e2e2a73fafd450f8a792e5c1143 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 21:26:56 +0200 Subject: [PATCH 03/20] feat: retry + circuit breaker + optional fallback backend for LLM providers --- .env.example | 7 ++ docs/architecture.md | 3 +- localrag/llm/factory.py | 22 +++++- localrag/llm/resilience.py | 108 +++++++++++++++++++++++++++++ localrag/settings.py | 7 ++ pyproject.toml | 2 + tests/test_llm_providers.py | 6 +- tests/test_llm_resilience.py | 130 +++++++++++++++++++++++++++++++++++ uv.lock | 13 ++++ 9 files changed, 293 insertions(+), 5 deletions(-) create mode 100644 localrag/llm/resilience.py create mode 100644 tests/test_llm_resilience.py diff --git a/.env.example b/.env.example index b094793..0012fa1 100644 --- a/.env.example +++ b/.env.example @@ -43,6 +43,13 @@ API_KEY= # LLM backend: ollama | openai | anthropic LLM_BACKEND=ollama +# Optional automatic failover backend when the primary trips its circuit breaker +# (ollama | openai | anthropic; empty disables fallback). +LLM_FALLBACK_BACKEND= +LLM_RETRY_MAX_ATTEMPTS=3 +LLM_CIRCUIT_FAIL_MAX=5 +LLM_CIRCUIT_RESET_TIMEOUT_SECONDS=30.0 + # Canonical embedding model alias (overrides OLLAMA_EMBED_MODEL when set). EMBEDDING_MODEL= diff --git a/docs/architecture.md b/docs/architecture.md index d9cf4f7..8b5e732 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -83,7 +83,8 @@ flowchart LR | `localrag/llm/providers/ollama.py` | Ollama HTTP provider (default, local) | | `localrag/llm/providers/openai_provider.py` | OpenAI chat completions | | `localrag/llm/providers/anthropic_provider.py` | Anthropic messages API | -| `localrag/llm/factory.py` | `build_provider(settings)` — selects provider by `LLM_BACKEND` env var | +| `localrag/llm/resilience.py` | `ResilientProvider` — retry-with-backoff (tenacity) + circuit breaker (pybreaker) wrapping any `BaseLLMProvider`; optional fallback provider on sustained failure | +| `localrag/llm/factory.py` | `build_provider(settings)` — selects provider by `LLM_BACKEND` env var; always returns a `ResilientProvider`-wrapped instance | | `localrag/llm/costs.py` | `estimate_cost_usd(model, tokens)` with prefix-match price table | ## Agent layer diff --git a/localrag/llm/factory.py b/localrag/llm/factory.py index db9714d..6691d04 100644 --- a/localrag/llm/factory.py +++ b/localrag/llm/factory.py @@ -6,12 +6,30 @@ from localrag.llm.providers.base import BaseLLMProvider from localrag.llm.providers.ollama import OllamaProvider from localrag.llm.providers.openai_provider import OpenAIProvider +from localrag.llm.resilience import ResilientProvider from localrag.settings import Settings def build_provider(settings: Settings) -> BaseLLMProvider: - """Return the provider selected by ``LLM_BACKEND``.""" - backend = settings.llm_backend.lower() + """Return the provider selected by ``LLM_BACKEND``, wrapped with retry + circuit breaker.""" + primary = _build_raw_provider(settings, settings.llm_backend) + fallback = ( + _build_raw_provider(settings, settings.llm_fallback_backend) + if settings.llm_fallback_backend + else None + ) + return ResilientProvider( + primary, + max_attempts=settings.llm_retry_max_attempts, + fail_max=settings.llm_circuit_fail_max, + reset_timeout_seconds=settings.llm_circuit_reset_timeout_seconds, + fallback_provider=fallback, + ) + + +def _build_raw_provider(settings: Settings, backend: str) -> BaseLLMProvider: + """Return the unwrapped provider for ``backend`` (no retry/circuit-breaker).""" + backend = backend.lower() if backend == "openai": return OpenAIProvider( api_key=settings.openai_api_key, diff --git a/localrag/llm/resilience.py b/localrag/llm/resilience.py new file mode 100644 index 0000000..dbfc576 --- /dev/null +++ b/localrag/llm/resilience.py @@ -0,0 +1,108 @@ +"""Retry-with-backoff + circuit-breaker wrapping for any BaseLLMProvider. + +Composes two independent concerns: transient failures get retried a few +times with exponential backoff; sustained failure trips the breaker open so +further calls fail fast (or fall back to a secondary provider) instead of +retry-storming a down backend. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Generator +from typing import Any, TypeVar + +import anthropic +import httpx +import openai +import pybreaker +from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential + +from localrag.llm.providers.base import BaseLLMProvider +from localrag.llm.types import LLMResponse + +logger = logging.getLogger(__name__) + +_T = TypeVar("_T") + +_RETRYABLE_EXCEPTION_TYPES: tuple[type[BaseException], ...] = ( + httpx.TransportError, + openai.APIConnectionError, + openai.APITimeoutError, + openai.RateLimitError, + anthropic.APIConnectionError, + anthropic.APITimeoutError, + anthropic.RateLimitError, +) + + +def _should_retry(exc: BaseException) -> bool: + if isinstance(exc, _RETRYABLE_EXCEPTION_TYPES): + return True + return isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code >= 500 + + +class ResilientProvider(BaseLLMProvider): + """Wraps a provider with retry-with-backoff and a circuit breaker.""" + + def __init__( + self, + provider: BaseLLMProvider, + *, + max_attempts: int = 3, + fail_max: int = 5, + reset_timeout_seconds: float = 30.0, + fallback_provider: BaseLLMProvider | None = None, + ) -> None: + self._provider = provider + self._fallback_provider = fallback_provider + self._breaker = pybreaker.CircuitBreaker( + fail_max=fail_max, reset_timeout=reset_timeout_seconds + ) + self._max_attempts = max_attempts + + def _retrying(self, func: Callable[[], _T]) -> _T: + wrapped = retry( + retry=retry_if_exception(_should_retry), + stop=stop_after_attempt(self._max_attempts), + wait=wait_exponential(multiplier=0.5, min=0.5, max=8), + reraise=True, + )(func) + return wrapped() + + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + try: + return self._breaker.call( + self._retrying, lambda: self._provider.generate(prompt, context, 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(prompt, context, model=model) + raise + + def stream( + self, prompt: str, context: list[str], *, model: str | None = 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) + 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) + return + raise + if first_event is not None: + yield first_event + yield from gen + + def count_tokens(self, text: str) -> int: + return self._provider.count_tokens(text) diff --git a/localrag/settings.py b/localrag/settings.py index fbeffaa..c5b9b83 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -117,6 +117,13 @@ class Settings(BaseSettings): # Agent settings (uses Anthropic tool-use) agent_model: str = "claude-haiku-4-5" + # Optional automatic failover backend when the primary trips its circuit breaker + # ("ollama" | "openai" | "anthropic"); empty disables fallback. + llm_fallback_backend: str = "" + llm_retry_max_attempts: int = 3 + llm_circuit_fail_max: int = 5 + llm_circuit_reset_timeout_seconds: float = 30.0 + log_level: str = "INFO" diff --git a/pyproject.toml b/pyproject.toml index 84ec15e..793ed29 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,8 @@ dependencies = [ "ragas>=0.4.3", "rank-bm25>=0.2.2", "python-multipart>=0.0.20", + "tenacity>=9", + "pybreaker>=1", ] [project.scripts] diff --git a/tests/test_llm_providers.py b/tests/test_llm_providers.py index 87958ed..b16cc68 100644 --- a/tests/test_llm_providers.py +++ b/tests/test_llm_providers.py @@ -14,6 +14,7 @@ from localrag.llm.providers.base import BaseLLMProvider from localrag.llm.providers.ollama import OllamaProvider from localrag.llm.providers.openai_provider import OpenAIProvider +from localrag.llm.resilience import ResilientProvider from localrag.llm.types import LLMResponse from localrag.settings import Settings @@ -91,7 +92,7 @@ def test_ollama_provider_cost_is_zero() -> None: ("ollama", OllamaProvider), ], ) -def test_build_provider_returns_correct_type( +def test_build_provider_wraps_correct_type_in_resilient_provider( backend: str, expected_type: type[BaseLLMProvider] ) -> None: settings = Settings( @@ -100,4 +101,5 @@ def test_build_provider_returns_correct_type( anthropic_api_key="sk-ant-test", ) provider = build_provider(settings) - assert isinstance(provider, expected_type) + assert isinstance(provider, ResilientProvider) + assert isinstance(provider._provider, expected_type) # noqa: SLF001 diff --git a/tests/test_llm_resilience.py b/tests/test_llm_resilience.py new file mode 100644 index 0000000..5fdc852 --- /dev/null +++ b/tests/test_llm_resilience.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +from collections.abc import Generator +from dataclasses import dataclass, field +from typing import Any + +import httpx +import pybreaker +import pytest + +from localrag.llm.providers.base import BaseLLMProvider +from localrag.llm.resilience import ResilientProvider +from localrag.llm.types import LLMResponse + + +@dataclass +class FlakyProvider(BaseLLMProvider): + """Fails ``fail_times`` times with a retryable error, then succeeds.""" + + fail_times: int + calls: list[str] = field(default_factory=list) + + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + self.calls.append(prompt) + if len(self.calls) <= self.fail_times: + request = httpx.Request("POST", "http://x/api/chat") + raise httpx.ConnectError("boom", request=request) + return LLMResponse( + answer="ok", model="m", tokens_used=1, latency_ms=0.0, estimated_cost_usd=0.0 + ) + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + yield {"type": "token", "token": "ok"} + yield {"type": "final", "sources": []} + + def count_tokens(self, text: str) -> int: + return len(text.split()) + + +@dataclass +class AlwaysFailsProvider(BaseLLMProvider): + calls: list[str] = field(default_factory=list) + + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + self.calls.append(prompt) + request = httpx.Request("POST", "http://x/api/chat") + raise httpx.ConnectError("always down", request=request) + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError + + def count_tokens(self, text: str) -> int: + return len(text.split()) + + +@dataclass +class FallbackProvider(BaseLLMProvider): + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + return LLMResponse( + answer="fallback", model="fb", tokens_used=1, latency_ms=0.0, estimated_cost_usd=0.0 + ) + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + yield {"type": "token", "token": "fallback"} + yield {"type": "final", "sources": []} + + def count_tokens(self, text: str) -> int: + return len(text.split()) + + +def test_resilient_provider_retries_transient_failure_then_succeeds() -> None: + flaky = FlakyProvider(fail_times=2) + resilient = ResilientProvider(flaky, max_attempts=3, fail_max=5, reset_timeout_seconds=30) + + response = resilient.generate("q", []) + + assert response.answer == "ok" + assert len(flaky.calls) == 3 + + +def test_resilient_provider_opens_circuit_after_fail_max_and_falls_back() -> None: + # pybreaker (>=1) raises CircuitBreakerError directly from the call that + # crosses fail_max, rather than propagating the underlying exception on + # that final call — so with fail_max=2, the first call surfaces the + # provider's own error; the second is where the breaker trips, and since a + # fallback is configured, ResilientProvider catches that and serves the + # fallback answer immediately instead of raising. + always_fails = AlwaysFailsProvider() + fallback = FallbackProvider() + resilient = ResilientProvider( + always_fails, + max_attempts=1, + fail_max=2, + reset_timeout_seconds=30, + fallback_provider=fallback, + ) + + with pytest.raises(httpx.ConnectError): + resilient.generate("q", []) + + # Circuit is now open — further calls go straight to the fallback, no further + # attempts against the primary provider. + response = resilient.generate("q", []) + assert response.answer == "fallback" + assert len(always_fails.calls) == 2 + + response_again = resilient.generate("q", []) + assert response_again.answer == "fallback" + assert len(always_fails.calls) == 2 + + +def test_resilient_provider_raises_when_circuit_open_and_no_fallback() -> None: + # With fail_max=1, the very first failure already crosses the threshold, so + # pybreaker raises CircuitBreakerError on that first call too (see note above). + always_fails = AlwaysFailsProvider() + resilient = ResilientProvider( + always_fails, max_attempts=1, fail_max=1, reset_timeout_seconds=30 + ) + + with pytest.raises(pybreaker.CircuitBreakerError): + resilient.generate("q", []) + + with pytest.raises(pybreaker.CircuitBreakerError): + resilient.generate("q", []) diff --git a/uv.lock b/uv.lock index d31c983..beabcee 100644 --- a/uv.lock +++ b/uv.lock @@ -1371,6 +1371,7 @@ dependencies = [ { name = "httpx" }, { name = "openai" }, { name = "prometheus-client" }, + { name = "pybreaker" }, { name = "pydantic-settings" }, { name = "pypdf" }, { name = "pypdfium2" }, @@ -1382,6 +1383,7 @@ dependencies = [ { name = "rich" }, { name = "sse-starlette" }, { name = "structlog" }, + { name = "tenacity" }, { name = "typer" }, { name = "uvicorn", extra = ["standard"] }, ] @@ -1407,6 +1409,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.28" }, { name = "openai", specifier = ">=2.33.0" }, { name = "prometheus-client", specifier = ">=0.25.0" }, + { name = "pybreaker", specifier = ">=1" }, { name = "pydantic-settings", specifier = ">=2.8" }, { name = "pypdf", specifier = ">=5.4" }, { name = "pypdfium2", specifier = ">=4.30" }, @@ -1418,6 +1421,7 @@ requires-dist = [ { name = "rich", specifier = ">=13" }, { name = "sse-starlette", specifier = ">=2.2" }, { name = "structlog", specifier = ">=25.5.0" }, + { name = "tenacity", specifier = ">=9" }, { name = "typer", specifier = ">=0.15" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.34" }, ] @@ -2466,6 +2470,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3b/30/4a675864877397179b09b720ee5fcb1cf772cf7bebc831989aff0a5f79c1/pybase64-1.4.3-cp314-cp314t-win_arm64.whl", hash = "sha256:e906aa08d4331e799400829e0f5e4177e76a3281e8a4bc82ba114c6b30e405c9", size = 31904, upload-time = "2025-12-06T13:25:26.826Z" }, ] +[[package]] +name = "pybreaker" +version = "1.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f2/89/fbf98e383f1ec6d117af2cd983efdb3eb7018b63834c427025764194cac2/pybreaker-1.4.1.tar.gz", hash = "sha256:8df2d245c73ba40c8242c56ffb4f12138fbadc23e296224740c2028ea9dc1178", size = 15555, upload-time = "2025-09-21T15:12:04.499Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/75/e64d3d40a741e2be21d69154f4e5c43a66f0c603c5ef11f49e01429a5932/pybreaker-1.4.1-py3-none-any.whl", hash = "sha256:b4dab4a05195b7f2a64a6c1a6c4ba7a96534ef56ea7210e6bcb59f28897160e0", size = 12915, upload-time = "2025-09-21T15:12:02.284Z" }, +] + [[package]] name = "pydantic" version = "2.13.4" From 969ee2effb5d0acf05aee576e204cf73d720f797 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 21:36:17 +0200 Subject: [PATCH 04/20] feat: metadata pre-filtering across vector and BM25 retrieval paths --- docs/rag-retrieval.md | 29 +++++++++++++++++++ localrag/api/schemas.py | 10 +++++++ localrag/api/service.py | 2 ++ localrag/rag/engine.py | 20 ++++++++++--- localrag/rag/retriever.py | 24 +++++++++++++--- localrag/storage/vector_store.py | 7 +++-- tests/test_api.py | 9 ++++-- tests/test_api_main_exceptions.py | 8 +++++- tests/test_freshness_decay.py | 6 ++-- tests/test_rag_engine.py | 16 ++++++++--- tests/test_retriever.py | 30 +++++++++++++++++-- tests/test_retriever_hybrid.py | 48 +++++++++++++++++++++++++++++-- tests/test_vector_store.py | 18 ++++++++++++ 13 files changed, 204 insertions(+), 23 deletions(-) diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index c58b2f7..ac3890e 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -55,6 +55,35 @@ deliberately disjoint with zero overlap: boundary-awareness (never splitting mid-table, mid-code-block, or mid-heading-section) substitutes for overlap there. This is an intentional design choice, not an oversight. +## Metadata pre-filtering + +`POST /query`'s `QueryRequest.metadata_filter` (`localrag/api/schemas.py`) accepts +an optional equality-only `dict[str, str]` filter applied to chunk metadata +**before** ranking, e.g.: + +```json +{"question": "...", "metadata_filter": {"source": "/docs/handbook.pdf"}} +``` + +It threads through `RAGEngine.answer` / `RAGEngine.stream_answer` +(`localrag/rag/engine.py`) into `Retriever.retrieve`'s `metadata_filter` +parameter (`localrag/rag/retriever.py`), which applies it on **both** retrieval +paths: + +1. **Vector search** — passed natively as Chroma's `where=` clause via + `VectorStore.query(embedding, top_k, where=...)` + (`localrag/storage/vector_store.py`), so Chroma itself excludes + non-matching chunks before the HNSW search returns results. +2. **BM25 search** (hybrid mode only) — applied client-side as an equality + check against each BM25 hit's metadata via the `_matches_filter` helper in + `localrag/rag/retriever.py`, since `rank_bm25` has no native filter concept. + +This is **equality-only** — it is not a full Chroma `$and`/`$or`/`$in` query +DSL. Every key/value pair in `metadata_filter` must match exactly +(`metadata.get(key) == value`) for a chunk to survive filtering; there is no +support for ranges, negation, or boolean combinators. Pairs naturally with the +`source`/`file_type` fields already written on every chunk during ingestion. + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: diff --git a/localrag/api/schemas.py b/localrag/api/schemas.py index e490aaf..d514e7d 100644 --- a/localrag/api/schemas.py +++ b/localrag/api/schemas.py @@ -58,6 +58,16 @@ class QueryRequest(BaseModel): ), examples=[5], ) + metadata_filter: dict[str, str] | None = Field( + default=None, + description=( + "Optional equality filters applied to chunk metadata before ranking, e.g. " + '{"source": "/docs/handbook.pdf"}. Applied natively to the vector search ' + "via Chroma's `where` clause and, for parity, as an equality check against " + "BM25 candidate metadata in hybrid mode." + ), + examples=[{"source": "/docs/handbook.pdf"}], + ) class IngestFileRequest(BaseModel): diff --git a/localrag/api/service.py b/localrag/api/service.py index de1a737..387779d 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -280,6 +280,7 @@ def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: contexts = engine.retriever.retrieve( question=request.question, n_results=request.n_results, + metadata_filter=request.metadata_filter, ) except RetrievalError as exc: raise RagApiError(int(exc.status_code), exc.detail) from exc @@ -325,6 +326,7 @@ def get_query_contexts(request: QueryRequest, engine: RAGEngine) -> list[dict[st return engine.retriever.retrieve( question=request.question, n_results=request.n_results, + metadata_filter=request.metadata_filter, ) except RetrievalError as exc: raise RagApiError(int(exc.status_code), exc.detail) from exc diff --git a/localrag/rag/engine.py b/localrag/rag/engine.py index c953b5d..c07a769 100644 --- a/localrag/rag/engine.py +++ b/localrag/rag/engine.py @@ -27,11 +27,17 @@ class RAGEngine: timeout_seconds: float = 180.0 def answer( - self, question: str, model: str | None = None, n_results: int | None = None + self, + question: str, + model: str | None = None, + n_results: int | None = None, + metadata_filter: dict[str, Any] | None = None, ) -> dict[str, object]: chunks: list[str] = [] sources: list[dict[str, object]] = [] - for event in self.stream_answer(question=question, model=model, n_results=n_results): + for event in self.stream_answer( + question=question, model=model, n_results=n_results, metadata_filter=metadata_filter + ): if event["type"] == "token": chunks.append(str(event["token"])) if event["type"] == "final": @@ -39,7 +45,11 @@ def answer( return {"answer": "".join(chunks).strip(), "sources": sources} def stream_answer( - self, question: str, model: str | None = None, n_results: int | None = None + self, + question: str, + model: str | None = None, + n_results: int | None = None, + metadata_filter: dict[str, Any] | None = None, ) -> Generator[dict[str, Any]]: logger.info( "rag_stream_start question_chars=%s model=%s n_results=%s", @@ -47,7 +57,9 @@ def stream_answer( model, n_results, ) - contexts = self.retriever.retrieve(question=question, n_results=n_results) + contexts = self.retriever.retrieve( + question=question, n_results=n_results, metadata_filter=metadata_filter + ) return self.stream_chat_from_contexts(contexts=contexts, question=question, model=model) def stream_chat_from_contexts( diff --git a/localrag/rag/retriever.py b/localrag/rag/retriever.py index a59a501..f37fb9f 100644 --- a/localrag/rag/retriever.py +++ b/localrag/rag/retriever.py @@ -17,6 +17,12 @@ logger = logging.getLogger(__name__) +def _matches_filter(metadata: dict[str, Any], metadata_filter: dict[str, Any] | None) -> bool: + if not metadata_filter: + return True + return all(metadata.get(key) == value for key, value in metadata_filter.items()) + + @dataclass class Retriever: settings: Settings @@ -24,7 +30,12 @@ class Retriever: vector_store: VectorStore bm25_index: Bm25Index | None = None - def retrieve(self, question: str, n_results: int | None = None) -> list[dict[str, Any]]: + def retrieve( + self, + question: str, + n_results: int | None = None, + metadata_filter: dict[str, Any] | None = None, + ) -> list[dict[str, Any]]: top_k = n_results if n_results is not None else self.settings.rag_top_k logger.debug("retrieve_embed_question top_k=%s question_chars=%s", top_k, len(question)) try: @@ -43,7 +54,9 @@ def retrieve(self, question: str, n_results: int | None = None) -> list[dict[str logger.error("retrieve_embed_invalid_response error=%s", exc) raise RetrievalError(HTTPStatus.BAD_GATEWAY, str(exc)) from exc - vector_hits = self._retrieve_vector_hits(embedding=embedding, top_k=max(top_k * 2, top_k)) + vector_hits = self._retrieve_vector_hits( + embedding=embedding, top_k=max(top_k * 2, top_k), where=metadata_filter + ) if self.settings.retrieval_mode != "hybrid" or self.bm25_index is None: return self.apply_freshness(vector_hits[:top_k]) @@ -57,14 +70,17 @@ def retrieve(self, question: str, n_results: int | None = None) -> list[dict[str "metadata": hit.metadata, } for hit in self.bm25_index.query(question, top_k=max(top_k * 2, top_k)) + if _matches_filter(hit.metadata, metadata_filter) ] return self.apply_freshness( self._fuse_results(vector_hits=vector_hits, bm25_hits=bm25_hits, top_k=top_k) ) - def _retrieve_vector_hits(self, embedding: list[float], top_k: int) -> list[dict[str, Any]]: + def _retrieve_vector_hits( + self, embedding: list[float], top_k: int, where: dict[str, Any] | None = None + ) -> list[dict[str, Any]]: try: - query_result = self.vector_store.query(embedding=embedding, top_k=top_k) + query_result = self.vector_store.query(embedding=embedding, top_k=top_k, where=where) except Exception as exc: logger.exception("retrieve_vector_store_query_failed top_k=%s", top_k) raise RetrievalError( diff --git a/localrag/storage/vector_store.py b/localrag/storage/vector_store.py index e8e38c0..7023c98 100644 --- a/localrag/storage/vector_store.py +++ b/localrag/storage/vector_store.py @@ -69,11 +69,14 @@ def delete_by_source(self, source: str) -> None: self.collection.delete(where={"source": source}) logger.debug("vector_delete_by_source source=%s", source) - def query(self, embedding: list[float], top_k: int) -> dict[str, Any]: - logger.debug("vector_query top_k=%s", top_k) + def query( + self, embedding: list[float], top_k: int, where: dict[str, Any] | None = None + ) -> dict[str, Any]: + logger.debug("vector_query top_k=%s where=%s", top_k, where) return self.collection.query( # type: ignore[return-value] query_embeddings=[embedding], # type: ignore[arg-type] n_results=top_k, + where=where, include=["documents", "metadatas", "distances"], ) diff --git a/tests/test_api.py b/tests/test_api.py index 9b4f511..d73a5a4 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -22,8 +22,13 @@ @dataclass class StubRetriever: - def retrieve(self, question: str, n_results: int | None = None) -> list[dict[str, Any]]: - _ = (question, n_results) + def retrieve( + self, + question: str, + n_results: int | None = None, + metadata_filter: dict[str, Any] | None = None, + ) -> list[dict[str, Any]]: + _ = (question, n_results, metadata_filter) return _STUB_CONTEXTS diff --git a/tests/test_api_main_exceptions.py b/tests/test_api_main_exceptions.py index 813cce4..c556689 100644 --- a/tests/test_api_main_exceptions.py +++ b/tests/test_api_main_exceptions.py @@ -59,7 +59,13 @@ def list_collection_names(self) -> list[str]: @dataclass class FailingRetriever: - def retrieve(self, question: str, n_results: int | None = None) -> list[object]: + def retrieve( + self, + question: str, + n_results: int | None = None, + metadata_filter: dict[str, object] | None = None, + ) -> list[object]: + _ = (question, n_results, metadata_filter) raise RetrievalError(HTTPStatus.BAD_GATEWAY, "Embedding service unavailable.") diff --git a/tests/test_freshness_decay.py b/tests/test_freshness_decay.py index 13669bc..1a984b2 100644 --- a/tests/test_freshness_decay.py +++ b/tests/test_freshness_decay.py @@ -52,8 +52,10 @@ def test_freshness_decay_prefers_recent_chunk() -> None: class FreshnessStore: @staticmethod - def query(embedding: list[float], top_k: int) -> dict[str, object]: - _ = (embedding, top_k) + def query( + embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k, where) return { "documents": [["old policy", "new policy"]], "metadatas": [ diff --git a/tests/test_rag_engine.py b/tests/test_rag_engine.py index f474efb..a33657c 100644 --- a/tests/test_rag_engine.py +++ b/tests/test_rag_engine.py @@ -15,8 +15,13 @@ class StubRetriever: contexts: list[dict[str, object]] - def retrieve(self, question: str, n_results: int | None = None) -> list[dict[str, object]]: - _ = (question, n_results) + def retrieve( + self, + question: str, + n_results: int | None = None, + metadata_filter: dict[str, object] | None = None, + ) -> list[dict[str, object]]: + _ = (question, n_results, metadata_filter) return self.contexts @@ -65,9 +70,12 @@ def test_rag_engine_answer_concatenates_tokens_and_returns_sources( engine = RAGEngine(settings=settings, retriever=StubRetriever(contexts=[])) # type: ignore[arg-type] def fake_stream_answer( - question: str, model: str | None = None, n_results: int | None = None + question: str, + model: str | None = None, + n_results: int | None = None, + metadata_filter: dict[str, object] | None = None, ) -> Generator[dict[str, object]]: - _ = (question, model, n_results) + _ = (question, model, n_results, metadata_filter) yield {"type": "token", "token": "Hello "} yield {"type": "token", "token": "World"} yield { diff --git a/tests/test_retriever.py b/tests/test_retriever.py index f131b71..9ea6bf1 100644 --- a/tests/test_retriever.py +++ b/tests/test_retriever.py @@ -21,8 +21,10 @@ def embed_text(self, text: str, *, model: str | None = None) -> list[float]: @dataclass class StubStore: - def query(self, embedding: list[float], top_k: int) -> dict[str, object]: - _ = (embedding, top_k) + def query( + self, embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k, where) return { "documents": [["chunk-a"]], "metadatas": [[{"source": "foo.md", "chunk_index": 0}]], @@ -70,6 +72,30 @@ def test_retriever_raises_retrieval_failure_when_ollama_embed_fails() -> None: assert excinfo.value.status_code == HTTPStatus.BAD_GATEWAY +def test_retriever_threads_metadata_filter_to_vector_store_where() -> None: + captured: dict[str, object] = {} + + @dataclass + class CapturingStore: + @staticmethod + def query( + embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k) + captured["where"] = where + return {"documents": [[]], "metadatas": [[]], "distances": [[]]} + + retriever = Retriever( + settings=Settings(), + embedder=StubEmbedder(), # type: ignore[arg-type] + vector_store=CapturingStore(), # type: ignore[arg-type] + ) + + retriever.retrieve("q", metadata_filter={"source": "a.md"}) + + assert captured["where"] == {"source": "a.md"} + + def test_retriever_raises_retrieval_failure_when_vector_query_fails() -> None: @dataclass class ExplodingStore: diff --git a/tests/test_retriever_hybrid.py b/tests/test_retriever_hybrid.py index f2e72fe..b8add02 100644 --- a/tests/test_retriever_hybrid.py +++ b/tests/test_retriever_hybrid.py @@ -18,8 +18,10 @@ def embed_text(text: str, *, model: str | None = None) -> list[float]: @dataclass class StubStore: @staticmethod - def query(embedding: list[float], top_k: int) -> dict[str, object]: - _ = (embedding, top_k) + def query( + embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k, where) return { "documents": [["nearby vector text", "exact token text"]], "metadatas": [ @@ -59,3 +61,45 @@ def test_retriever_hybrid_fuses_vector_and_bm25() -> None: contexts = retriever.retrieve("ERR_QUIC_PROTOCOL_ERROR", n_results=2) assert contexts[0]["source"] == "exact.md" + + +def test_retriever_hybrid_applies_metadata_filter_to_bm25_hits() -> None: + @dataclass + class FilteringStore: + @staticmethod + def query( + embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k) + documents = ["nearby vector text", "exact token text"] + metadatas = [ + {"source": "nearby.md", "chunk_index": 0}, + {"source": "exact.md", "chunk_index": 1}, + ] + if where: + keep = [ + i + for i, md in enumerate(metadatas) + if all(md.get(k) == v for k, v in where.items()) + ] + documents = [documents[i] for i in keep] + metadatas = [metadatas[i] for i in keep] + return { + "documents": [documents], + "metadatas": [metadatas], + "distances": [[0.01] * len(documents)], + } + + settings = Settings(retrieval_mode="hybrid", rrf_k=1) + retriever = Retriever( + settings=settings, + embedder=StubEmbedder(), # type: ignore[arg-type] + vector_store=FilteringStore(), # type: ignore[arg-type] + bm25_index=StubBm25Index(), # type: ignore[arg-type] + ) + + contexts = retriever.retrieve( + "ERR_QUIC_PROTOCOL_ERROR", n_results=2, metadata_filter={"source": "nearby.md"} + ) + + assert all(context["source"] == "nearby.md" for context in contexts) diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 5da02da..8db84f7 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -52,12 +52,14 @@ def query( query_embeddings: list[list[float]], n_results: int, include: list[str], + where: dict[str, object] | None = None, ) -> dict[str, object]: self.query_calls.append( { "query_embeddings": query_embeddings, "n_results": n_results, "include": include, + "where": where, } ) return self.query_result @@ -182,6 +184,22 @@ def test_vector_store_upsert_delete_query_and_list_collections() -> None: assert client.deleted_collections == ["col-1"] +def test_vector_store_query_passes_where_filter_to_collection() -> None: + collection = FakeCollection( + upsert_calls=[], + delete_calls=[], + query_calls=[], + query_result={"documents": [[]], "metadatas": [[]], "distances": [[]]}, + get_return={}, + ) + client = FakeClient(collections=[], deleted_collections=[]) + store = VectorStore(client=client, collection=collection) # type: ignore[arg-type] + + store.query(embedding=[0.1], top_k=3, where={"source": "a.md"}) + + assert collection.query_calls[0]["where"] == {"source": "a.md"} # type: ignore[index] + + def test_vector_store_list_distinct_sources() -> None: collection = FakeCollection( upsert_calls=[], From 7ba7117b7d73cb28e6ca06165bcfdc15d3523761 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 22:10:58 +0200 Subject: [PATCH 05/20] feat: expand top retrieval hits to full heading section before prompting --- .env.example | 2 ++ docs/rag-retrieval.md | 20 +++++++++++ localrag/rag/prompt.py | 2 +- localrag/rag/retriever.py | 30 +++++++++++++--- localrag/settings.py | 4 +++ localrag/storage/vector_store.py | 18 ++++++++++ tests/test_rag_prompt.py | 15 ++++++++ tests/test_retriever.py | 62 ++++++++++++++++++++++++++++++++ tests/test_vector_store.py | 22 ++++++++++++ 9 files changed, 170 insertions(+), 5 deletions(-) diff --git a/.env.example b/.env.example index 0012fa1..8fd7470 100644 --- a/.env.example +++ b/.env.example @@ -31,6 +31,8 @@ RETRIEVAL_MODE=hybrid BM25_WEIGHT=0.5 RRF_K=60 FRESHNESS_HALF_LIFE_DAYS=30.0 +# Expand top retrieval hits to their full markdown section (via heading_path) before prompting. Set false to disable. +PARENT_EXPANSION_ENABLED=true RAG_SYSTEM_PROMPT=You are a helpful assistant. Answer only based on the provided context. API_HOST=0.0.0.0 diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index ac3890e..9717b1e 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -84,6 +84,26 @@ DSL. Every key/value pair in `metadata_filter` must match exactly support for ranges, negation, or boolean combinators. Pairs naturally with the `source`/`file_type` fields already written on every chunk during ingestion. +## Parent-section expansion + +After ranking, fusion, and freshness decay, `Retriever._expand_to_parent_section` +(`localrag/rag/retriever.py`) expands top retrieval hits that carry a non-empty +`heading_path` chunk metadata value to the **full sibling-chunk section** they +belong to, via `VectorStore.get_chunks_by_heading(source, heading_path)` +(`localrag/storage/vector_store.py`), which fetches every chunk sharing that +`source` + `heading_path` pair and returns them sorted by `chunk_index`. The +merged section text is joined with `"\n\n"` and stored on the context dict as +`expanded_text`, while the originally matched chunk's `text` (and +`chunk_index`) are retained unchanged so citations (`SourceRef`) still point at +the precise matched chunk. `localrag/rag/prompt.py::build_prompt` prefers +`expanded_text` over `text` when present when composing the LLM prompt, so the +model sees the whole section instead of just the one matching sentence. + +Controlled by `PARENT_EXPANSION_ENABLED` (default `true`); set to `false` to +skip expansion and prompt with only the originally matched chunk text. Hits +with an empty `heading_path`, or whose section has only a single chunk, are +left unexpanded. + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: diff --git a/localrag/rag/prompt.py b/localrag/rag/prompt.py index b0eb855..9559f99 100644 --- a/localrag/rag/prompt.py +++ b/localrag/rag/prompt.py @@ -6,7 +6,7 @@ def build_prompt(system_prompt: str, question: str, contexts: list[dict[str, obj for index, context in enumerate(contexts, start=1): source = context.get("source", "unknown") chunk_index = context.get("chunk_index", -1) - text = context.get("text", "") + text = context.get("expanded_text") or context.get("text", "") context_blocks.append(f"[{index}] source={source} chunk={chunk_index}\n{text}") joined_context = "\n\n".join(context_blocks) if context_blocks else "No context found." diff --git a/localrag/rag/retriever.py b/localrag/rag/retriever.py index f37fb9f..3a11dac 100644 --- a/localrag/rag/retriever.py +++ b/localrag/rag/retriever.py @@ -58,7 +58,7 @@ def retrieve( embedding=embedding, top_k=max(top_k * 2, top_k), where=metadata_filter ) if self.settings.retrieval_mode != "hybrid" or self.bm25_index is None: - return self.apply_freshness(vector_hits[:top_k]) + return self._expand_to_parent_section(self.apply_freshness(vector_hits[:top_k])) bm25_hits = [ { @@ -72,9 +72,8 @@ def retrieve( for hit in self.bm25_index.query(question, top_k=max(top_k * 2, top_k)) if _matches_filter(hit.metadata, metadata_filter) ] - return self.apply_freshness( - self._fuse_results(vector_hits=vector_hits, bm25_hits=bm25_hits, top_k=top_k) - ) + fused = self._fuse_results(vector_hits=vector_hits, bm25_hits=bm25_hits, top_k=top_k) + return self._expand_to_parent_section(self.apply_freshness(fused)) def _retrieve_vector_hits( self, embedding: list[float], top_k: int, where: dict[str, Any] | None = None @@ -181,3 +180,26 @@ def _hit_key(hit: dict[str, Any]) -> tuple[str, int]: source = str(hit.get("source", "unknown")) chunk_index = int(hit.get("chunk_index", -1)) return source, chunk_index + + def _expand_to_parent_section(self, contexts: list[dict[str, Any]]) -> list[dict[str, Any]]: + if not self.settings.parent_expansion_enabled: + return contexts + expanded: list[dict[str, Any]] = [] + for context in contexts: + metadata = context.get("metadata") or {} + heading_path = metadata.get("heading_path", "") + source = context.get("source", "unknown") + if not heading_path: + expanded.append(context) + continue + siblings = self.vector_store.get_chunks_by_heading( + source=str(source), heading_path=str(heading_path) + ) + if len(siblings) <= 1: + expanded.append(context) + continue + merged_text = "\n\n".join(text for _, text in siblings) + new_context = dict(context) + new_context["expanded_text"] = merged_text + expanded.append(new_context) + return expanded diff --git a/localrag/settings.py b/localrag/settings.py index c5b9b83..b9971b7 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -45,6 +45,9 @@ class Settings(BaseSettings): **RAG** — ``rag_top_k`` is how many chunks are retrieved for context. ``rag_system_prompt`` is the system message for the answering model. + When ``parent_expansion_enabled`` is true (default), top hits with a + non-empty ``heading_path`` are expanded to their full sibling-chunk + section before prompting; set false to disable. **API** — ``api_host`` / ``api_port`` are the uvicorn bind address and port. @@ -89,6 +92,7 @@ class Settings(BaseSettings): bm25_weight: float = 0.5 rrf_k: int = 60 freshness_half_life_days: float = 30.0 + parent_expansion_enabled: bool = True rag_system_prompt: str = ( "You are a helpful assistant. Answer only based on the provided context." ) diff --git a/localrag/storage/vector_store.py b/localrag/storage/vector_store.py index 7023c98..b67207b 100644 --- a/localrag/storage/vector_store.py +++ b/localrag/storage/vector_store.py @@ -80,6 +80,24 @@ def query( include=["documents", "metadatas", "distances"], ) + def get_chunks_by_heading(self, source: str, heading_path: str) -> list[tuple[int, str]]: + raw = self.collection.get( + where={"$and": [{"source": source}, {"heading_path": heading_path}]}, + include=["documents", "metadatas"], + ) + documents = raw.get("documents") or [] + metadatas = raw.get("metadatas") or [] + pairs: list[tuple[int, str]] = [] + for document, metadata in zip(documents, metadatas, strict=False): + if not isinstance(document, str): + continue + metadata_map = metadata if isinstance(metadata, dict) else {} + chunk_index = metadata_map.get("chunk_index") + if isinstance(chunk_index, int): + pairs.append((chunk_index, document)) + pairs.sort(key=lambda pair: pair[0]) + return pairs + def list_distinct_sources(self) -> list[str]: raw = self.collection.get(include=["metadatas"]) metadatas = raw.get("metadatas") diff --git a/tests/test_rag_prompt.py b/tests/test_rag_prompt.py index c2549db..edf3894 100644 --- a/tests/test_rag_prompt.py +++ b/tests/test_rag_prompt.py @@ -21,3 +21,18 @@ def test_build_prompt_includes_context_blocks() -> None: assert "[1] source=foo.md chunk=2\nhello" in out assert "[2] source=bar.md chunk=0\nworld" in out assert "Question:\nQ" in out + + +def test_build_prompt_prefers_expanded_text_over_matched_text() -> None: + contexts = [ + { + "source": "guide.md", + "chunk_index": 0, + "text": "matched sentence only", + "expanded_text": "full section text including matched sentence", + } + ] + out = build_prompt(system_prompt="SYS", question="Q", contexts=contexts) + + assert "full section text including matched sentence" in out + assert "matched sentence only" not in out diff --git a/tests/test_retriever.py b/tests/test_retriever.py index 9ea6bf1..d779f5c 100644 --- a/tests/test_retriever.py +++ b/tests/test_retriever.py @@ -112,3 +112,65 @@ def query(self, embedding: list[float], top_k: int) -> dict[str, object]: retriever.retrieve("q") assert excinfo.value.status_code == HTTPStatus.SERVICE_UNAVAILABLE + + +def test_retriever_expands_top_hits_to_full_heading_section() -> None: + @dataclass + class ExpandableStore: + @staticmethod + def query( + embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k, where) + return { + "documents": [["Section intro sentence."]], + "metadatas": [[{"source": "guide.md", "chunk_index": 0, "heading_path": "Setup"}]], + "distances": [[0.05]], + } + + @staticmethod + def get_chunks_by_heading(source: str, heading_path: str) -> list[tuple[int, str]]: + assert source == "guide.md" + assert heading_path == "Setup" + return [ + (0, "Section intro sentence."), + (1, "Second sentence with the install command."), + ] + + retriever = Retriever( + settings=Settings(), + embedder=StubEmbedder(), # type: ignore[arg-type] + vector_store=ExpandableStore(), # type: ignore[arg-type] + ) + + contexts = retriever.retrieve("how do I install this") + + assert contexts[0]["expanded_text"] == ( + "Section intro sentence.\n\nSecond sentence with the install command." + ) + assert contexts[0]["text"] == "Section intro sentence." + + +def test_retriever_skips_expansion_when_heading_path_empty() -> None: + @dataclass + class FlatStore: + @staticmethod + def query( + embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k, where) + return { + "documents": [["plain text chunk"]], + "metadatas": [[{"source": "notes.txt", "chunk_index": 0, "heading_path": ""}]], + "distances": [[0.05]], + } + + retriever = Retriever( + settings=Settings(), + embedder=StubEmbedder(), # type: ignore[arg-type] + vector_store=FlatStore(), # type: ignore[arg-type] + ) + + contexts = retriever.retrieve("q") + + assert "expanded_text" not in contexts[0] diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 8db84f7..432ad3d 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -241,6 +241,28 @@ def test_vector_store_get_all_chunks_returns_id_doc_metadata_triplets() -> None: ] +def test_vector_store_get_chunks_by_heading_returns_sorted_pairs() -> None: + collection = FakeCollection( + upsert_calls=[], + delete_calls=[], + query_calls=[], + query_result={}, + get_return={ + "documents": ["second", "first"], + "metadatas": [ + {"source": "guide.md", "chunk_index": 1, "heading_path": "Setup"}, + {"source": "guide.md", "chunk_index": 0, "heading_path": "Setup"}, + ], + }, + ) + client = FakeClient(collections=[], deleted_collections=[]) + store = VectorStore(client=client, collection=collection) # type: ignore[arg-type] + + pairs = store.get_chunks_by_heading(source="guide.md", heading_path="Setup") + + assert pairs == [(0, "first"), (1, "second")] + + def test_vector_store_create_initializes_persistent_client( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: From bc231b6f45432617c3aa9f9787844eca2314a55c Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 22:13:34 +0200 Subject: [PATCH 06/20] fix: route RAGEngine through the LLM provider abstraction (LLM_BACKEND now actually governs /query) --- docs/architecture.md | 10 +- localrag/api/dependencies.py | 5 +- localrag/llm/providers/anthropic_provider.py | 32 ++++++ localrag/llm/providers/base.py | 10 ++ localrag/llm/providers/ollama.py | 51 +++++++++ localrag/llm/providers/openai_provider.py | 26 +++++ localrag/llm/resilience.py | 31 ++++++ localrag/rag/engine.py | 56 ++-------- tests/test_llm_providers.py | 11 ++ tests/test_llm_resilience.py | 24 +++++ tests/test_rag_engine.py | 107 ++++++++++++------- 11 files changed, 272 insertions(+), 91 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 8b5e732..a41e614 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -44,8 +44,8 @@ flowchart LR - **Ingest:** files → `loader` / `ingestion/parsers/*` → text → structural chunker (`localrag/ingestion/structural_chunker.py`) or fixed fallback (`localrag/ingestion/chunker.py`) → `OllamaEmbedder` → `VectorStore` (Chroma, persistent path from settings). The **HTTP** ingest flow runs path decode, existence checks, and `INGEST_ROOTS` in `localrag/api/service.py` (`ingest_file` / `ingest_directory`), then calls `IngestionService`; optional per-request `embed_model` overrides `OLLAMA_EMBED_MODEL`. Failures raise `IngestApiError` → JSON in `main.py`. CLI ingests call `IngestionService` directly. `POST /ingest/upload` (`ingest_upload` in `service.py`) takes a multipart file instead of a server path: it validates the extension against `loader.SUPPORTED_EXTENSIONS`, streams it to disk under `UPLOAD_DIR` in 1 MiB chunks while enforcing `UPLOAD_MAX_BYTES` (bypassing `INGEST_ROOTS`, since the server picks the destination), then calls `IngestionService.ingest_file` the same way. See the endpoint's OpenAPI description for upload limitations (no AV scan, extension-only validation, single file per request). - **Rebuild:** `POST /collections/rebuild` and `localrag collections rebuild` list distinct `source` values in the active collection, drop vectors for missing files, and re-chunk/re-embed remaining paths (optional `embed_model` override). Implemented in `IngestionService.rebuild_collection`. -- **Query (JSON):** `POST /query` returns a complete `QueryResponse` (answer, sources, latency_ms, model) from `query_json` in `localrag/api/service.py`. Retrieval supports vector-only and hybrid (vector + BM25 with reciprocal-rank fusion), then applies optional freshness decay based on chunk `ingested_at`. Requires `X-API-Key` when `API_KEY` is set. -- **Query (SSE stream):** `POST /query/stream` streams tokens as Server-Sent Events. Retrieval runs synchronously first (`get_query_contexts`) so errors map to HTTP before SSE starts, then tokens are mapped via `iter_query_sse_events`. +- **Query (JSON):** `POST /query` returns a complete `QueryResponse` (answer, sources, latency_ms, model) from `query_json` in `localrag/api/service.py`. Retrieval supports vector-only and hybrid (vector + BM25 with reciprocal-rank fusion), then applies optional freshness decay based on chunk `ingested_at`. `RAGEngine` generates the answer via its injected `provider` (a `BaseLLMProvider` built by `llm/factory.py::build_provider`, resilience-wrapped), so `LLM_BACKEND` genuinely governs which backend answers `/query` — it is no longer hard-wired to Ollama. Requires `X-API-Key` when `API_KEY` is set. +- **Query (SSE stream):** `POST /query/stream` streams tokens as Server-Sent Events. Retrieval runs synchronously first (`get_query_contexts`) so errors map to HTTP before SSE starts, then tokens are mapped via `iter_query_sse_events`. Token streaming likewise goes through `RAGEngine.provider.stream_from_prompt(...)`, so `LLM_BACKEND` governs the streaming path too. - **Metrics:** `GET /metrics` exposes Prometheus metrics via `prometheus_client` (router at `localrag/api/routers/metrics.py`). No auth required. ## Package map @@ -78,15 +78,17 @@ flowchart LR | Path | Role | | --- | --- | -| `localrag/llm/providers/base.py` | `BaseLLMProvider` ABC with `generate(prompt, context)` and `stream(...)` | +| `localrag/llm/providers/base.py` | `BaseLLMProvider` ABC with `generate(prompt, context)` / `stream(...)` (context-list contract for direct scripting use) and `generate_from_prompt(prompt)` / `stream_from_prompt(prompt)` (already-built-prompt contract used by `RAGEngine`) | | `localrag/llm/types.py` | `LLMResponse` dataclass (answer, model, tokens_used, latency_ms, estimated_cost_usd) | | `localrag/llm/providers/ollama.py` | Ollama HTTP provider (default, local) | | `localrag/llm/providers/openai_provider.py` | OpenAI chat completions | | `localrag/llm/providers/anthropic_provider.py` | Anthropic messages API | -| `localrag/llm/resilience.py` | `ResilientProvider` — retry-with-backoff (tenacity) + circuit breaker (pybreaker) wrapping any `BaseLLMProvider`; optional fallback provider on sustained failure | +| `localrag/llm/resilience.py` | `ResilientProvider` — retry-with-backoff (tenacity) + circuit breaker (pybreaker) wrapping any `BaseLLMProvider` (including the `*_from_prompt` methods); optional fallback provider on sustained failure | | `localrag/llm/factory.py` | `build_provider(settings)` — selects provider by `LLM_BACKEND` env var; always returns a `ResilientProvider`-wrapped instance | | `localrag/llm/costs.py` | `estimate_cost_usd(model, tokens)` with prefix-match price table | +`RAGEngine` (`localrag/rag/engine.py`) holds a `provider: BaseLLMProvider` field (injected in `localrag/api/dependencies.py::get_engine` via `build_provider(settings)`) and builds its own citation-rich prompt with `localrag.rag.prompt.build_prompt`, then calls `self.provider.stream_from_prompt(prompt, model=model)` — it no longer talks to Ollama's `/api/chat` directly, so switching `LLM_BACKEND` to `openai` or `anthropic` actually changes what answers `/query` and `/query/stream`. + ## Agent layer `localrag/agent/service.py` exposes `run_agent(question, engine, api_key, model)`: diff --git a/localrag/api/dependencies.py b/localrag/api/dependencies.py index ffd3133..b523c3c 100644 --- a/localrag/api/dependencies.py +++ b/localrag/api/dependencies.py @@ -9,6 +9,7 @@ from localrag.api.repository import ChromaCollectionRepository from localrag.ingestion.embedder import OllamaEmbedder from localrag.ingestion.service import IngestionService +from localrag.llm.factory import build_provider from localrag.rag.bm25_index import Bm25Index from localrag.rag.engine import RAGEngine from localrag.rag.retriever import Retriever @@ -50,7 +51,9 @@ def get_retriever() -> Retriever: @lru_cache(maxsize=1) def get_engine() -> RAGEngine: settings = get_settings() - return RAGEngine(settings=settings, retriever=get_retriever()) + return RAGEngine( + settings=settings, retriever=get_retriever(), provider=build_provider(settings) + ) @lru_cache(maxsize=1) diff --git a/localrag/llm/providers/anthropic_provider.py b/localrag/llm/providers/anthropic_provider.py index 6f419c5..49f74c0 100644 --- a/localrag/llm/providers/anthropic_provider.py +++ b/localrag/llm/providers/anthropic_provider.py @@ -76,5 +76,37 @@ def stream( yield {"type": "token", "token": text} yield {"type": "final", "sources": []} + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + t0 = time.perf_counter() + used_model = model or self._default_model + resp = self._client.messages.create( + model=used_model, + max_tokens=_MAX_TOKENS, + messages=[{"role": "user", "content": prompt}], + ) + latency_ms = (time.perf_counter() - t0) * 1000 + answer = next((b.text for b in resp.content if hasattr(b, "text")), "") # type: ignore[union-attr] + tokens = resp.usage.input_tokens + resp.usage.output_tokens + return LLMResponse( + answer=answer.strip(), + model=used_model, + tokens_used=tokens, + latency_ms=latency_ms, + estimated_cost_usd=estimate_cost_usd(used_model, tokens), + ) + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + used_model = model or self._default_model + with self._client.messages.stream( + model=used_model, + max_tokens=_MAX_TOKENS, + messages=[{"role": "user", "content": prompt}], + ) as stream: + for text in stream.text_stream: + yield {"type": "token", "token": text} + yield {"type": "final", "sources": []} + def count_tokens(self, text: str) -> int: return len(text.split()) diff --git a/localrag/llm/providers/base.py b/localrag/llm/providers/base.py index ef95748..c56300a 100644 --- a/localrag/llm/providers/base.py +++ b/localrag/llm/providers/base.py @@ -32,6 +32,16 @@ def stream( ) -> Generator[dict[str, Any]]: """Yield token/final events matching the existing RAG engine contract.""" + @abstractmethod + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + """Return a complete response for an already-fully-built prompt string.""" + + @abstractmethod + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + """Yield token/final events for an already-fully-built prompt string.""" + @abstractmethod def count_tokens(self, text: str) -> int: """Approximate token count for ``text``.""" diff --git a/localrag/llm/providers/ollama.py b/localrag/llm/providers/ollama.py index 144ad2f..3972ed5 100644 --- a/localrag/llm/providers/ollama.py +++ b/localrag/llm/providers/ollama.py @@ -99,5 +99,56 @@ def stream( break yield {"type": "final", "sources": []} + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + t0 = time.perf_counter() + chunks: list[str] = [] + for event in self.stream_from_prompt(prompt, model=model): + if event["type"] == "token": + chunks.append(str(event["token"])) + latency_ms = (time.perf_counter() - t0) * 1000 + used_model = model or self._default_model + answer = "".join(chunks).strip() + tokens = len(answer.split()) + return LLMResponse( + answer=answer, + model=used_model, + tokens_used=tokens, + latency_ms=latency_ms, + estimated_cost_usd=estimate_cost_usd("_default_ollama", tokens), + ) + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + used_model = model or self._default_model + chat_request = OllamaChatRequest( + model=used_model, + messages=[OllamaChatMessage(role="user", content=prompt)], + stream=True, + ) + with ( + httpx.Client(timeout=self._timeout) as client, + client.stream( + "POST", + f"{self._base_url}/api/chat", + json=chat_request.model_dump(mode="json", exclude_none=True), + ) as resp, + ): + resp.raise_for_status() + for line in resp.iter_lines(): + if not line: + continue + try: + chunk = parse_ollama_json_line(line, OllamaChatStreamChunk) + except ValueError: + continue + msg = chunk.message + token = msg.content if msg and msg.content else "" + if token: + yield {"type": "token", "token": token} + if chunk.done: + break + yield {"type": "final", "sources": []} + def count_tokens(self, text: str) -> int: return len(text.split()) diff --git a/localrag/llm/providers/openai_provider.py b/localrag/llm/providers/openai_provider.py index 5173aba..74d5541 100644 --- a/localrag/llm/providers/openai_provider.py +++ b/localrag/llm/providers/openai_provider.py @@ -70,5 +70,31 @@ def stream( yield {"type": "token", "token": text} yield {"type": "final", "sources": []} + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + t0 = time.perf_counter() + used_model = model or self._default_model + messages = [{"role": "user", "content": prompt}] + resp = self._client.chat.completions.create(model=used_model, messages=messages) # type: ignore[arg-type] + latency_ms = (time.perf_counter() - t0) * 1000 + answer = resp.choices[0].message.content or "" + tokens = resp.usage.total_tokens if resp.usage else self.count_tokens(answer) + return LLMResponse( + answer=answer.strip(), + model=used_model, + tokens_used=tokens, + latency_ms=latency_ms, + estimated_cost_usd=estimate_cost_usd(used_model, tokens), + ) + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + used_model = model or self._default_model + messages = [{"role": "user", "content": prompt}] + with self._client.chat.completions.stream(model=used_model, messages=messages) as stream: # type: ignore[arg-type] + for text in stream.text_stream: # type: ignore[attr-defined] + yield {"type": "token", "token": text} + yield {"type": "final", "sources": []} + def count_tokens(self, text: str) -> int: return len(text.split()) diff --git a/localrag/llm/resilience.py b/localrag/llm/resilience.py index dbfc576..8c13efe 100644 --- a/localrag/llm/resilience.py +++ b/localrag/llm/resilience.py @@ -104,5 +104,36 @@ def start() -> tuple[Generator[dict[str, Any]], dict[str, Any] | None]: yield first_event yield from gen + 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 + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + def start() -> tuple[Generator[dict[str, Any]], dict[str, Any] | None]: + gen = self._provider.stream_from_prompt(prompt, model=model) + 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 + 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 c07a769..0808c02 100644 --- a/localrag/rag/engine.py +++ b/localrag/rag/engine.py @@ -5,14 +5,7 @@ from dataclasses import dataclass from typing import Any -import httpx - -from localrag.ollama.schemas import ( - OllamaChatMessage, - OllamaChatRequest, - OllamaChatStreamChunk, - parse_ollama_json_line, -) +from localrag.llm.providers.base import BaseLLMProvider from localrag.rag.prompt import build_prompt from localrag.rag.retriever import Retriever from localrag.settings import Settings @@ -24,7 +17,7 @@ class RAGEngine: settings: Settings retriever: Retriever - timeout_seconds: float = 180.0 + provider: BaseLLMProvider def answer( self, @@ -85,47 +78,10 @@ def _stream_chat_tokens( question=question, contexts=contexts, ) - runtime_model = model or self.settings.ollama_llm_model - chat_request = OllamaChatRequest( - model=runtime_model, - messages=[OllamaChatMessage(role="user", content=prompt)], - stream=True, - ) - - try: - with ( - httpx.Client(timeout=self.timeout_seconds) as client, - client.stream( - "POST", - f"{self.settings.ollama_base_url}/api/chat", - json=chat_request.model_dump(mode="json", exclude_none=True), - ) as resp, - ): - resp.raise_for_status() - for line in resp.iter_lines(): - if not line: - continue - try: - chunk = parse_ollama_json_line(line, OllamaChatStreamChunk) - except ValueError: - logger.warning("rag_ollama_bad_chunk line_chars=%s", len(line)) - continue - msg = chunk.message - token = msg.content if msg and msg.content else "" - if token: - yield {"type": "token", "token": token} - if chunk.done: - break - except httpx.HTTPError as exc: - logger.error( - "rag_ollama_chat_http_error url=%s model=%s error=%s", - self.settings.ollama_base_url, - runtime_model, - exc, - ) - raise - - logger.info("rag_stream_done model=%s", runtime_model) + for event in self.provider.stream_from_prompt(prompt, model=model): + if event["type"] == "token": + yield event + logger.info("rag_stream_done") yield {"type": "final", "sources": self._extract_sources(contexts)} @staticmethod diff --git a/tests/test_llm_providers.py b/tests/test_llm_providers.py index b16cc68..579bc65 100644 --- a/tests/test_llm_providers.py +++ b/tests/test_llm_providers.py @@ -33,6 +33,17 @@ def stream( yield {"type": "token", "token": "ok"} yield {"type": "final", "sources": []} + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + return LLMResponse( + answer="ok", model="test", tokens_used=1, latency_ms=0.0, estimated_cost_usd=0.0 + ) + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + yield {"type": "token", "token": "ok"} + yield {"type": "final", "sources": []} + def count_tokens(self, text: str) -> int: return len(text.split()) diff --git a/tests/test_llm_resilience.py b/tests/test_llm_resilience.py index 5fdc852..788d49c 100644 --- a/tests/test_llm_resilience.py +++ b/tests/test_llm_resilience.py @@ -35,6 +35,14 @@ def stream( yield {"type": "token", "token": "ok"} yield {"type": "final", "sources": []} + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + return self.generate(prompt, [], model=model) + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + yield from self.stream(prompt, [], model=model) + def count_tokens(self, text: str) -> int: return len(text.split()) @@ -53,6 +61,14 @@ def stream( ) -> Generator[dict[str, Any]]: raise NotImplementedError + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + return self.generate(prompt, [], model=model) + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + yield from self.stream(prompt, [], model=model) + def count_tokens(self, text: str) -> int: return len(text.split()) @@ -70,6 +86,14 @@ def stream( yield {"type": "token", "token": "fallback"} yield {"type": "final", "sources": []} + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + return self.generate(prompt, [], model=model) + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + yield from self.stream(prompt, [], model=model) + def count_tokens(self, text: str) -> int: return len(text.split()) diff --git a/tests/test_rag_engine.py b/tests/test_rag_engine.py index a33657c..82e660f 100644 --- a/tests/test_rag_engine.py +++ b/tests/test_rag_engine.py @@ -1,12 +1,13 @@ from __future__ import annotations from collections.abc import Generator -from dataclasses import dataclass +from dataclasses import dataclass, field +from typing import Any -import httpx import pytest -import respx +from localrag.llm.providers.base import BaseLLMProvider +from localrag.llm.types import LLMResponse from localrag.rag.engine import RAGEngine from localrag.settings import Settings @@ -25,32 +26,44 @@ def retrieve( return self.contexts -@respx.mock +@dataclass +class FakeProvider(BaseLLMProvider): + tokens: list[str] + prompts_seen: list[str] = field(default_factory=list) + + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + raise NotImplementedError + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError + + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + raise NotImplementedError + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + self.prompts_seen.append(prompt) + for token in self.tokens: + yield {"type": "token", "token": token} + yield {"type": "final", "sources": []} + + def count_tokens(self, text: str) -> int: + return len(text.split()) + + def test_rag_engine_stream_answer_yields_tokens_and_dedupes_sources() -> None: - settings = Settings( - ollama_base_url="http://ollama:11434", - rag_system_prompt="SYS", - ollama_llm_model="llm", - ) + settings = Settings(rag_system_prompt="SYS") contexts = [ {"text": "chunk-one", "source": "a.md", "chunk_index": 1}, # Duplicate (same source + chunk_index) to exercise dedupe. {"text": "chunk-one", "source": "a.md", "chunk_index": 1}, ] - engine = RAGEngine(settings=settings, retriever=StubRetriever(contexts=contexts)) - - url = f"{settings.ollama_base_url}/api/chat" - ndjson_lines = [ - '{"model":"llm","message":{"role":"assistant","content":"Hello"},"done":false}', - "not-json", - "", - '{"model":"llm","message":{"role":"assistant","content":" world"},"done":false}', - '{"model":"llm","message":{"role":"assistant","content":""},"done":true}', - ] - content = ("\n".join(ndjson_lines) + "\n").encode("utf-8") - - respx.post(url).mock( - return_value=httpx.Response(200, content=content), + provider = FakeProvider(tokens=["Hello", " world"]) + engine = RAGEngine( + settings=settings, retriever=StubRetriever(contexts=contexts), provider=provider ) events = list(engine.stream_answer(question="Q", model="llm", n_results=3)) @@ -61,13 +74,16 @@ def test_rag_engine_stream_answer_yields_tokens_and_dedupes_sources() -> None: final = events[-1] assert final["type"] == "final" assert final["sources"] == [{"source": "a.md", "chunk_index": 1}] + assert "SYS" in provider.prompts_seen[0] + assert "chunk-one" in provider.prompts_seen[0] def test_rag_engine_answer_concatenates_tokens_and_returns_sources( monkeypatch: pytest.MonkeyPatch, ) -> None: - settings = Settings(rag_system_prompt="SYS", ollama_base_url="http://ollama:11434") - engine = RAGEngine(settings=settings, retriever=StubRetriever(contexts=[])) # type: ignore[arg-type] + settings = Settings(rag_system_prompt="SYS") + provider = FakeProvider(tokens=[]) + engine = RAGEngine(settings=settings, retriever=StubRetriever(contexts=[]), provider=provider) def fake_stream_answer( question: str, @@ -93,17 +109,36 @@ def fake_stream_answer( ] -@respx.mock -def test_rag_engine_stream_answer_raises_on_http_error() -> None: - settings = Settings( - ollama_base_url="http://ollama:11434", - rag_system_prompt="SYS", - ollama_llm_model="llm", - ) - engine = RAGEngine(settings=settings, retriever=StubRetriever(contexts=[])) +def test_rag_engine_stream_answer_propagates_provider_error() -> None: + settings = Settings(rag_system_prompt="SYS") + + @dataclass + class ExplodingProvider(BaseLLMProvider): + def generate( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> LLMResponse: + raise NotImplementedError + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError - url = f"{settings.ollama_base_url}/api/chat" - respx.post(url).mock(return_value=httpx.Response(500, content=b"oops\n")) + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + raise NotImplementedError + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise RuntimeError("provider down") + yield # pragma: no cover — unreachable, keeps this a generator + + def count_tokens(self, text: str) -> int: + return len(text.split()) + + engine = RAGEngine( + settings=settings, retriever=StubRetriever(contexts=[]), provider=ExplodingProvider() + ) - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(RuntimeError, match="provider down"): list(engine.stream_answer(question="Q", model="llm", n_results=1)) From e3af91b04cb442ce0adceddf4a22923da616a443 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 06:44:23 +0200 Subject: [PATCH 07/20] feat: optional cross-encoder reranking; drop dead RRF weight branch --- .env.example | 5 + docs/rag-retrieval.md | 20 ++ localrag/api/dependencies.py | 10 + localrag/rag/reranker.py | 45 ++++ localrag/rag/retriever.py | 59 +++-- localrag/settings.py | 10 + pyproject.toml | 3 + tests/test_reranker.py | 33 +++ tests/test_retriever.py | 43 ++++ uv.lock | 460 +++++++++++++++++++++++++++++++++-- 10 files changed, 645 insertions(+), 43 deletions(-) create mode 100644 localrag/rag/reranker.py create mode 100644 tests/test_reranker.py diff --git a/.env.example b/.env.example index 8fd7470..1c13bb0 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,11 @@ FRESHNESS_HALF_LIFE_DAYS=30.0 PARENT_EXPANSION_ENABLED=true RAG_SYSTEM_PROMPT=You are a helpful assistant. Answer only based on the provided context. +# Optional cross-encoder reranking (requires `uv sync --extra rerank`). +RERANK_ENABLED=false +RERANK_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2 +RERANK_FETCH_K=20 + API_HOST=0.0.0.0 API_PORT=8000 diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index 9717b1e..770d5cf 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -104,6 +104,26 @@ skip expansion and prompt with only the originally matched chunk text. Hits with an empty `heading_path`, or whose section has only a single chunk, are left unexpanded. +## Cross-encoder reranking (optional) + +Disabled by default (`RERANK_ENABLED=false`). When enabled, `Retriever.retrieve` +(`localrag/rag/retriever.py`) over-fetches `RERANK_FETCH_K` candidates from the +vector/hybrid path instead of the default `top_k * 2`, and a local +`cross-encoder/ms-marco-MiniLM-L-6-v2` model (`RERANK_MODEL`, served via +`sentence-transformers` — install with `uv sync --extra rerank`) re-scores each +`(question, chunk_text)` pair through `CrossEncoderReranker.rerank` +(`localrag/rag/reranker.py`) and trims the candidate list down to `top_k`, +best-first, adding a `rerank_score` key to each context. + +Reranking runs on the raw fused/vector candidate list — **before** +`apply_freshness` and `_expand_to_parent_section` — so it acts as the final +relevance step, and freshness decay / parent-section expansion still apply to +the reranked, already-trimmed top-`k` results (matching how those two behave +when reranking is disabled). `localrag/api/dependencies.py::get_reranker` +builds the `CrossEncoderReranker` only when `RERANK_ENABLED=true`, mirroring +the pluggable-provider shape in `localrag/llm/factory.py` (nothing imports +`sentence-transformers` unless the feature is turned on). + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: diff --git a/localrag/api/dependencies.py b/localrag/api/dependencies.py index b523c3c..950b1fc 100644 --- a/localrag/api/dependencies.py +++ b/localrag/api/dependencies.py @@ -12,6 +12,7 @@ from localrag.llm.factory import build_provider from localrag.rag.bm25_index import Bm25Index from localrag.rag.engine import RAGEngine +from localrag.rag.reranker import CrossEncoderReranker from localrag.rag.retriever import Retriever from localrag.settings import Settings, get_settings from localrag.storage.vector_store import VectorStore @@ -37,6 +38,14 @@ def get_embedder() -> OllamaEmbedder: ) +@lru_cache(maxsize=1) +def get_reranker() -> CrossEncoderReranker | None: + settings = get_settings() + if not settings.rerank_enabled: + return None + return CrossEncoderReranker(model_name=settings.rerank_model) + + @lru_cache(maxsize=1) def get_retriever() -> Retriever: settings = get_settings() @@ -45,6 +54,7 @@ def get_retriever() -> Retriever: embedder=get_embedder(), vector_store=get_vector_store(), bm25_index=get_bm25_index(), + reranker=get_reranker(), ) diff --git a/localrag/rag/reranker.py b/localrag/rag/reranker.py new file mode 100644 index 0000000..7ac45ac --- /dev/null +++ b/localrag/rag/reranker.py @@ -0,0 +1,45 @@ +"""Optional cross-encoder reranking for retrieved contexts (not loaded unless enabled).""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Any, Protocol + +logger = logging.getLogger(__name__) + + +class CrossEncoderModel(Protocol): + def predict(self, pairs: list[tuple[str, str]]) -> list[float]: ... + + +def _load_default_model(model_name: str) -> CrossEncoderModel: + from sentence_transformers import CrossEncoder # noqa: PLC0415 — optional dependency + + return CrossEncoder(model_name) + + +@dataclass +class CrossEncoderReranker: + model_name: str + _model: CrossEncoderModel | None = field(default=None, repr=False) + + def _get_model(self) -> CrossEncoderModel: + if self._model is None: + self._model = _load_default_model(self.model_name) + return self._model + + def rerank( + self, question: str, contexts: list[dict[str, Any]], top_k: int + ) -> list[dict[str, Any]]: + if not contexts: + return contexts + pairs = [(question, str(context.get("text", ""))) for context in contexts] + scores = self._get_model().predict(pairs) + rescored = [ + {**context, "rerank_score": float(score)} + for context, score in zip(contexts, scores, strict=True) + ] + rescored.sort(key=lambda context: context["rerank_score"], reverse=True) + logger.debug("rerank_done candidates=%s top_k=%s", len(contexts), top_k) + return rescored[:top_k] diff --git a/localrag/rag/retriever.py b/localrag/rag/retriever.py index 3a11dac..234cf9d 100644 --- a/localrag/rag/retriever.py +++ b/localrag/rag/retriever.py @@ -11,6 +11,7 @@ from localrag.ingestion.embedder import OllamaEmbedder from localrag.rag.bm25_index import Bm25Index from localrag.rag.exceptions import RetrievalError +from localrag.rag.reranker import CrossEncoderReranker from localrag.settings import Settings from localrag.storage.vector_store import VectorStore @@ -29,6 +30,7 @@ class Retriever: embedder: OllamaEmbedder vector_store: VectorStore bm25_index: Bm25Index | None = None + reranker: CrossEncoderReranker | None = None def retrieve( self, @@ -37,6 +39,11 @@ def retrieve( metadata_filter: dict[str, Any] | None = None, ) -> list[dict[str, Any]]: top_k = n_results if n_results is not None else self.settings.rag_top_k + fetch_k = ( + max(self.settings.rerank_fetch_k, top_k) + if self.reranker is not None + else max(top_k * 2, top_k) + ) logger.debug("retrieve_embed_question top_k=%s question_chars=%s", top_k, len(question)) try: embedding = self.embedder.embed_text(question) @@ -55,25 +62,33 @@ def retrieve( raise RetrievalError(HTTPStatus.BAD_GATEWAY, str(exc)) from exc vector_hits = self._retrieve_vector_hits( - embedding=embedding, top_k=max(top_k * 2, top_k), where=metadata_filter + embedding=embedding, top_k=fetch_k, where=metadata_filter ) if self.settings.retrieval_mode != "hybrid" or self.bm25_index is None: - return self._expand_to_parent_section(self.apply_freshness(vector_hits[:top_k])) - - bm25_hits = [ - { - "text": hit.text, - "source": hit.metadata.get("source", "unknown"), - "chunk_index": hit.metadata.get("chunk_index", -1), - "score": hit.score, - "ingested_at": hit.metadata.get("ingested_at"), - "metadata": hit.metadata, - } - for hit in self.bm25_index.query(question, top_k=max(top_k * 2, top_k)) - if _matches_filter(hit.metadata, metadata_filter) - ] - fused = self._fuse_results(vector_hits=vector_hits, bm25_hits=bm25_hits, top_k=top_k) - return self._expand_to_parent_section(self.apply_freshness(fused)) + candidates = vector_hits + else: + bm25_hits = [ + { + "text": hit.text, + "source": hit.metadata.get("source", "unknown"), + "chunk_index": hit.metadata.get("chunk_index", -1), + "score": hit.score, + "ingested_at": hit.metadata.get("ingested_at"), + "metadata": hit.metadata, + } + for hit in self.bm25_index.query(question, top_k=fetch_k) + if _matches_filter(hit.metadata, metadata_filter) + ] + candidates = self._fuse_results( + vector_hits=vector_hits, bm25_hits=bm25_hits, top_k=fetch_k + ) + + if self.reranker is not None: + candidates = self.reranker.rerank(question, candidates, top_k=top_k) + else: + candidates = candidates[:top_k] + + return self._expand_to_parent_section(self.apply_freshness(candidates)) def _retrieve_vector_hits( self, embedding: list[float], top_k: int, where: dict[str, Any] | None = None @@ -126,17 +141,11 @@ def _fuse_results( for rank, hit in enumerate(vector_sorted, start=1): key = self._hit_key(hit) candidate_map[key] = hit - if self.settings.bm25_weight == 0.5: - score_map[key] = score_map.get(key, 0.0) + 1.0 / (rrf_k + rank) - else: - score_map[key] = score_map.get(key, 0.0) + vector_weight / (rrf_k + rank) + score_map[key] = score_map.get(key, 0.0) + vector_weight / (rrf_k + rank) for rank, hit in enumerate(bm25_sorted, start=1): key = self._hit_key(hit) candidate_map[key] = hit - if self.settings.bm25_weight == 0.5: - score_map[key] = score_map.get(key, 0.0) + 1.0 / (rrf_k + rank) - else: - score_map[key] = score_map.get(key, 0.0) + bm25_weight / (rrf_k + rank) + score_map[key] = score_map.get(key, 0.0) + bm25_weight / (rrf_k + rank) ranked_keys = sorted(score_map.keys(), key=lambda key: score_map[key], reverse=True)[:top_k] fused: list[dict[str, Any]] = [] diff --git a/localrag/settings.py b/localrag/settings.py index b9971b7..21152e8 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -49,6 +49,11 @@ class Settings(BaseSettings): non-empty ``heading_path`` are expanded to their full sibling-chunk section before prompting; set false to disable. + **Reranking** — When ``rerank_enabled`` is true (default false, requires + ``uv sync --extra rerank``), retrieval over-fetches ``rerank_fetch_k`` + candidates and a local cross-encoder (``rerank_model``) re-scores and + trims them to ``rag_top_k`` before freshness/expansion. + **API** — ``api_host`` / ``api_port`` are the uvicorn bind address and port. **Logging** — ``log_level`` is the minimum level for the ``localrag`` logger @@ -97,6 +102,11 @@ class Settings(BaseSettings): "You are a helpful assistant. Answer only based on the provided context." ) + # Optional cross-encoder reranking (requires `uv sync --extra rerank`). + rerank_enabled: bool = False + rerank_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2" + rerank_fetch_k: int = 20 + api_host: str = "0.0.0.0" # nosec B104 — configurable bind address, default intentional api_port: int = 8000 diff --git a/pyproject.toml b/pyproject.toml index 793ed29..274ca97 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,9 @@ dependencies = [ [project.scripts] localrag = "localrag.cli.app:app" +[project.optional-dependencies] +rerank = ["sentence-transformers>=3"] + [dependency-groups] dev = [ "ruff", diff --git a/tests/test_reranker.py b/tests/test_reranker.py new file mode 100644 index 0000000..0db406e --- /dev/null +++ b/tests/test_reranker.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from localrag.rag.reranker import CrossEncoderReranker + + +@dataclass +class FakeModel: + scores: list[float] + + def predict(self, pairs: list[tuple[str, str]]) -> list[float]: + assert len(pairs) == len(self.scores) + return self.scores + + +def test_cross_encoder_reranker_reorders_by_score_and_trims_to_top_k() -> None: + contexts = [ + {"text": "low relevance", "source": "a.md", "chunk_index": 0}, + {"text": "high relevance", "source": "b.md", "chunk_index": 0}, + {"text": "mid relevance", "source": "c.md", "chunk_index": 0}, + ] + reranker = CrossEncoderReranker(model_name="unused", _model=FakeModel(scores=[0.1, 0.9, 0.5])) + + out = reranker.rerank("question", contexts, top_k=2) + + assert [c["source"] for c in out] == ["b.md", "c.md"] + assert out[0]["rerank_score"] == 0.9 + + +def test_cross_encoder_reranker_returns_empty_list_unchanged() -> None: + reranker = CrossEncoderReranker(model_name="unused", _model=FakeModel(scores=[])) + assert reranker.rerank("q", [], top_k=5) == [] diff --git a/tests/test_retriever.py b/tests/test_retriever.py index d779f5c..bb43393 100644 --- a/tests/test_retriever.py +++ b/tests/test_retriever.py @@ -151,6 +151,49 @@ def get_chunks_by_heading(source: str, heading_path: str) -> list[tuple[int, str assert contexts[0]["text"] == "Section intro sentence." +def test_retriever_applies_reranker_over_widened_candidate_pool() -> None: + @dataclass + class TwoDocStore: + @staticmethod + def query( + embedding: list[float], top_k: int, where: dict[str, object] | None = None + ) -> dict[str, object]: + _ = (embedding, top_k, where) + return { + "documents": [["nearby vector text", "exact token text"]], + "metadatas": [ + [ + {"source": "nearby.md", "chunk_index": 0}, + {"source": "exact.md", "chunk_index": 1}, + ] + ], + "distances": [[0.01, 0.4]], + } + + @dataclass + class FakeReranker: + calls: list[tuple[str, int]] + + def rerank( + self, question: str, contexts: list[dict[str, object]], top_k: int + ) -> list[dict[str, object]]: + self.calls.append((question, top_k)) + return list(reversed(contexts))[:top_k] + + reranker = FakeReranker(calls=[]) + retriever = Retriever( + settings=Settings(retrieval_mode="vector", rerank_fetch_k=2), + embedder=StubEmbedder(), # type: ignore[arg-type] + vector_store=TwoDocStore(), # type: ignore[arg-type] + reranker=reranker, # type: ignore[arg-type] + ) + + contexts = retriever.retrieve("q", n_results=1) + + assert reranker.calls == [("q", 1)] + assert contexts[0]["source"] == "exact.md" + + def test_retriever_skips_expansion_when_heading_path_empty() -> None: @dataclass class FlatStore: diff --git a/uv.lock b/uv.lock index beabcee..0aee460 100644 --- a/uv.lock +++ b/uv.lock @@ -511,6 +511,70 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/61/e8/cb8e80d6f9f55b99588625062822bf946cf03ed06315df4bd8397f5632a1/coverage-7.14.0-py3-none-any.whl", hash = "sha256:8de5b61163aee3d05c8a2beab6f47913df7981dad1baf82c414d99158c286ab1", size = 211764, upload-time = "2026-05-10T18:02:29.538Z" }, ] +[[package]] +name = "cuda-bindings" +version = "13.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/cc/6e/2394f8163360f8391f8f1b7e72d300a82724edb81a7b7084c799fbd4c91f/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9efb21c1ee64981e184b9e0ba5eb3179e5ba3d4b51665a6cb52b8ef3d01a7cbf", size = 5920504, upload-time = "2026-05-29T23:11:56.883Z" }, + { url = "https://files.pythonhosted.org/packages/34/c2/ef9b6a63f7dc432712a462c816662e662e00d38caa9b861c8c2588195d03/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2732904099e0a4d4db774a5fc6d91ee95fae065b4d2ecabb4968c5fe2406c9d7", size = 6476660, upload-time = "2026-05-29T23:11:59.188Z" }, + { url = "https://files.pythonhosted.org/packages/b1/81/bff68ce829999c1e4209c761bbf903b1c06ec570416ddb25020864ad5907/cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ab2f74ed65bfef4163ba07a8db16f1085e0729291db12a2423aff84ee8278b8", size = 6013639, upload-time = "2026-05-29T23:12:03.509Z" }, + { url = "https://files.pythonhosted.org/packages/d4/e0/c8a1f0c8f9ffdea4f5fe6dbab89b326cef4d85caf489dad39e209da89416/cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:efd4c814d311ec08c981f6dded1dbe7d4b371067ee4f6c14cccec4bde9590f80", size = 6534419, upload-time = "2026-05-29T23:12:05.633Z" }, + { url = "https://files.pythonhosted.org/packages/52/b8/83b1f563925b290f2d11a01a77a84013ba56052fe3653a5bef3ccfbb43d6/cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3c772dfff49681541d59630c90f858e173ac926b9c593a2b7123f2a1043cc76", size = 5809771, upload-time = "2026-05-29T23:12:10.422Z" }, + { url = "https://files.pythonhosted.org/packages/12/20/e79b4bfe98f075195afb6343d41c498f9dbd2d161d7021d4d28bceb83581/cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:36febb7c1079d68a981dbbd8d5a67235b399802b82075c9388624719607e52b9", size = 6358584, upload-time = "2026-05-29T23:12:12.767Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.5.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/53/8fc9b0cdc5b7f62746e6a01b85b6461e5ae27f871010a5fcf8fa6950766d/cuda_pathfinder-1.5.6-py3-none-any.whl", hash = "sha256:7e4c07c117b78ba1fb35dac4c444d21f3677b1b1ff56175c53a8e3025c5b43c0", size = 52972, upload-time = "2026-06-30T00:58:04.34Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, +] + +[package.optional-dependencies] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "sys_platform == 'linux'" }, +] +cufile = [ + { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" }, +] +curand = [ + { name = "nvidia-curand", marker = "sys_platform == 'linux'" }, +] +cusolver = [ + { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" }, +] +nvtx = [ + { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, +] + [[package]] name = "dataclasses-json" version = "0.6.7" @@ -1054,6 +1118,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f9/8e/7def204fea9f9be8b3c21a6f2dd6c020cf56c7d5ff753e0e23ed7f9ea57e/jiter-0.13.0-cp314-cp314t-win_arm64.whl", hash = "sha256:2c26cf47e2cad140fa23b6d58d435a7c0161f5c514284802f25e87fddfe11024", size = 187152, upload-time = "2026-02-02T12:37:22.124Z" }, ] +[[package]] +name = "joblib" +version = "1.5.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/f2/d34e8b3a08a9cc79a50b2208a93dce981fe615b64d5a4d4abee421d898df/joblib-1.5.3.tar.gz", hash = "sha256:8561a3269e6801106863fd0d6d84bb737be9e7631e33aaed3fb9ce5953688da3", size = 331603, upload-time = "2025-12-15T08:41:46.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl", hash = "sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713", size = 309071, upload-time = "2025-12-15T08:41:44.973Z" }, +] + [[package]] name = "jsonpatch" version = "1.33" @@ -1388,6 +1461,11 @@ dependencies = [ { name = "uvicorn", extra = ["standard"] }, ] +[package.optional-dependencies] +rerank = [ + { name = "sentence-transformers" }, +] + [package.dev-dependencies] dev = [ { name = "bandit" }, @@ -1419,12 +1497,14 @@ requires-dist = [ { name = "ragas", specifier = ">=0.4.3" }, { name = "rank-bm25", specifier = ">=0.2.2" }, { name = "rich", specifier = ">=13" }, + { name = "sentence-transformers", marker = "extra == 'rerank'", specifier = ">=3" }, { name = "sse-starlette", specifier = ">=2.2" }, { name = "structlog", specifier = ">=25.5.0" }, { name = "tenacity", specifier = ">=9" }, { name = "typer", specifier = ">=0.15" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.34" }, ] +provides-extras = ["rerank"] [package.metadata.requires-dev] dev = [ @@ -1651,6 +1731,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/0f/59204bf136d1201f8d7884cfbaf7498c5b4674e87a4c693f9bde63741ce1/mmh3-5.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:dfd51b4c56b673dfbc43d7d27ef857dd91124801e2806c69bb45585ce0fa019b", size = 40391, upload-time = "2026-03-05T15:55:56.697Z" }, ] +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + [[package]] name = "multidict" version = "6.7.1" @@ -1795,6 +1884,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, ] +[[package]] +name = "narwhals" +version = "2.23.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e8/ac/66ed1fc6e38a0c0f330627ec5c5d597990d6159b6712b82af0ad2c65f06c/narwhals-2.23.0.tar.gz", hash = "sha256:13e7ff5b4bb4a2f77b907c2e4d8a76e273dfc1323a3c997440a2f9fd26aed408", size = 656209, upload-time = "2026-07-01T11:21:53.278Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/4e/afc8c31605cb8be1d3bb4438c4d979daa104dab6306cd2b87abe9c3a7299/narwhals-2.23.0-py3-none-any.whl", hash = "sha256:769e7b9ab102c93d8fa019f6b4cd1a657909b04a20bf6210e5a35aae06814ae9", size = 458938, upload-time = "2026-07-01T11:21:51.677Z" }, +] + [[package]] name = "nest-asyncio" version = "1.6.0" @@ -1872,6 +1970,158 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/bb/e1c71a4295b1b1d1393d50dbb4f2a36283c6859d9d3892e84f00ec5a91d5/numpy-2.4.6-cp314-cp314t-win_arm64.whl", hash = "sha256:0c9136e14ed34a9e343a31c533d78a9813a69a3148332bce5e9821cb2f996e66", size = 10565867, upload-time = "2026-05-18T23:36:47.114Z" }, ] +[[package]] +name = "nvidia-cublas" +version = "13.1.1.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, + { url = "https://files.pythonhosted.org/packages/3b/cd/154ca20c38269e05eff77c1464e6c1da89f50a6390b565e9d82e06bc11e1/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436", size = 423138758, upload-time = "2026-04-08T18:46:58.655Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, + { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, + { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime" +version = "13.0.96" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, + { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu13" +version = "9.20.0.48" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, + { url = "https://files.pythonhosted.org/packages/6e/5e/edb9c0ae051602c3ccaffe424256463636d639e27d7f302dde9975ef9e7a/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304", size = 366173588, upload-time = "2026-03-09T19:29:34.474Z" }, +] + +[[package]] +name = "nvidia-cufft" +version = "12.0.0.61" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, +] + +[[package]] +name = "nvidia-cufile" +version = "1.15.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, + { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, +] + +[[package]] +name = "nvidia-curand" +version = "10.4.0.35" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, +] + +[[package]] +name = "nvidia-cusolver" +version = "12.0.4.66" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cusparse", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, + { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, +] + +[[package]] +name = "nvidia-cusparse" +version = "12.6.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu13" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" }, + { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" }, +] + +[[package]] +name = "nvidia-nccl-cu13" +version = "2.29.7" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/72/0d/daf50d44177ee0cbc7ff0a0c91eb5ff676c82be42f9a970bc7597f440c3a/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5", size = 206014712, upload-time = "2026-03-03T05:34:20.843Z" }, + { url = "https://files.pythonhosted.org/packages/67/f4/58e4e91b6919367c7aafb8e36fce9aad1a3047e536bf7e2fd560927d3a4c/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d", size = 205976000, upload-time = "2026-03-03T05:36:24.472Z" }, +] + +[[package]] +name = "nvidia-nvjitlink" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b", size = 40713933, upload-time = "2025-09-04T08:35:43.553Z" }, + { url = "https://files.pythonhosted.org/packages/ab/2c/93c5250e64df4f894f1cbb397c6fd71f79813f9fd79d7cd61de3f97b3c2d/nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c", size = 38768748, upload-time = "2025-09-04T08:35:20.008Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu13" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, + { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, +] + +[[package]] +name = "nvidia-nvtx" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, + { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, +] + [[package]] name = "oauthlib" version = "3.3.1" @@ -3059,6 +3309,63 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/d5/bc97ff895ec35cf3925d4bd60f3b39d822f377a446906ec9bcc87405e59b/ruff-0.15.14-py3-none-win_arm64.whl", hash = "sha256:ff47b90a9ef6a40c9e2f3b479c1fb78531adf055b94c1eba0a7ba04b31951826", size = 11208607, upload-time = "2026-05-21T14:34:26.525Z" }, ] +[[package]] +name = "safetensors" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, +] + +[[package]] +name = "scikit-learn" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "joblib" }, + { name = "narwhals" }, + { name = "numpy" }, + { name = "scipy" }, + { name = "threadpoolctl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fa/6f/37092bdb25f712817231799fc5674d8e704066a8a70c1d2d40517e18b4ab/scikit_learn-1.9.0.tar.gz", hash = "sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557", size = 7750767, upload-time = "2026-06-02T11:54:32.706Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/01/cf3310626b6d48d3e9be69a1223f9180360b5e6edb045f50fade723ce494/scikit_learn-1.9.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:80746d63bd4b6eaca54d36fe5feaf4d28bb38dc6f9470f81c7cad7c40155f119", size = 8705188, upload-time = "2026-06-02T11:53:41.964Z" }, + { url = "https://files.pythonhosted.org/packages/3e/04/5acd7ae280c5f93b6ac5ef6cdec14eef4c8d1cd91d85b3292989c94d96b1/scikit_learn-1.9.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5b934c45c252844a91d69fda3a34cff5e7307e1db10d77cb10a3980312c74713", size = 8228299, upload-time = "2026-06-02T11:53:44.817Z" }, + { url = "https://files.pythonhosted.org/packages/0c/39/ffe829a5b8ecb40a518724a997794657fdc354ada5e8fe8e64d998c0bac9/scikit_learn-1.9.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:38c3dcb9a1ffb85505ec53d54c7b4aea0cff70050425a7760c2af661ac85df05", size = 8789690, upload-time = "2026-06-02T11:53:47.461Z" }, + { url = "https://files.pythonhosted.org/packages/1f/88/8dab5de10c638c083772a6be83a3d8106ced492f74a928c8693638e5bb50/scikit_learn-1.9.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:da76d09304a4706db7cc1e3ebaa3b6b98a67365cc11d2996c4f1e58ba47df714", size = 9087723, upload-time = "2026-06-02T11:53:50.702Z" }, + { url = "https://files.pythonhosted.org/packages/20/3f/7917ca72464038f6240ec70c29f94862d08a34a74291ae4d4ec5eb8186a0/scikit_learn-1.9.0-cp313-cp313-win_amd64.whl", hash = "sha256:5808d98f15c6bf6d9d96d2348c1997392a5888ce7097e664105f930c4bca1277", size = 8184330, upload-time = "2026-06-02T11:53:53.396Z" }, + { url = "https://files.pythonhosted.org/packages/78/c7/15739eb2f61fda3c54639e9942414e5a19ad8a8d1f5a3266afad7cb7df80/scikit_learn-1.9.0-cp313-cp313-win_arm64.whl", hash = "sha256:d77f54c017633791bc0225a43e2f8d03745fdcfe4880268fcc4df15f505dec2e", size = 7840653, upload-time = "2026-06-02T11:53:56.035Z" }, + { url = "https://files.pythonhosted.org/packages/f4/7d/c9a35cf59b20a86fec24d306f1547b78dec194b08d367ce2a3e4854169d9/scikit_learn-1.9.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:9656acd4e93f74e0b66c8a36c88830a99252dfa900044d36bc2212ae89a47162", size = 8713289, upload-time = "2026-06-02T11:53:58.788Z" }, + { url = "https://files.pythonhosted.org/packages/3c/a7/552a7821597c632b907f7bfe8f36f9f572777af8ef8a48353041cf8e091a/scikit_learn-1.9.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:24360002ae845e7866522b0a5bbf690802e7bc388cac8663502e78aa98598aa2", size = 8245141, upload-time = "2026-06-02T11:54:01.694Z" }, + { url = "https://files.pythonhosted.org/packages/7d/79/f4a0c4fe9711154cddabf913471153af79056382ddc612cfe5ee0ff4b72e/scikit_learn-1.9.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5162ad10a418c8a282dde04c9aa06965de3e9a65f33c1440c0ae69bb1a09d913", size = 8847671, upload-time = "2026-06-02T11:54:04.448Z" }, + { url = "https://files.pythonhosted.org/packages/f0/af/4d72d9e475ac83719160c662619e4bf7b95c19507cd582e7d0167a3c3dae/scikit_learn-1.9.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fea2cc5677ab49d6f5bade978c866da44957b712d92e9635e8b4f723013c3cb", size = 9118104, upload-time = "2026-06-02T11:54:07.205Z" }, + { url = "https://files.pythonhosted.org/packages/a2/d5/6a58eea2cb9abbb9b3f2bb8b2cfb3243d1152d69f442d256c7af71304769/scikit_learn-1.9.0-cp314-cp314-win_amd64.whl", hash = "sha256:64fa347efc1c839c487433e40c5144d38c336e8a2b59c81aa8660373945c2673", size = 8290674, upload-time = "2026-06-02T11:54:10.087Z" }, + { url = "https://files.pythonhosted.org/packages/65/5b/d4c879cf358f1187141cf90ced473f087183489090244f50c124a2ee478b/scikit_learn-1.9.0-cp314-cp314-win_arm64.whl", hash = "sha256:1b944b6db288f6b926e3650026ddafb988929de95d11fc2cc5fa117773c9ba42", size = 7978807, upload-time = "2026-06-02T11:54:12.769Z" }, + { url = "https://files.pythonhosted.org/packages/8a/43/bfae3121ec67ae09150d453c442c7c1cc166e9aefe056e6ab3b7728a5cfc/scikit_learn-1.9.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:4ccacf04ca5f4b492158a5f28afe0ace43f81b2571e4b9a66d34848b46128949", size = 9031941, upload-time = "2026-06-02T11:54:15.436Z" }, + { url = "https://files.pythonhosted.org/packages/75/b0/20a4546eb17f3b25d3c66df15810411c14ed5065bcfab50b53c96fb627b2/scikit_learn-1.9.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:ee1a8db2c18c08e34c7412d4b10be1cac214cd4ea7dc9715a6a327eb49a37c96", size = 8613528, upload-time = "2026-06-02T11:54:18.842Z" }, + { url = "https://files.pythonhosted.org/packages/18/3c/e440e039bb82cd19004edaaad00acbde0fb9b461083c3ecf37941c557312/scikit_learn-1.9.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:147e9329ef0e39f75d4cffa02b2aa48d827832684926cd5210d9a2cb5c57246b", size = 8855050, upload-time = "2026-06-02T11:54:21.699Z" }, + { url = "https://files.pythonhosted.org/packages/43/26/b341b8dab5998da6270a3a42c2152c578501354d36f944b5856757035ef8/scikit_learn-1.9.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5bad8f8b9950321b54c965fdcbac6c6c55e79e16646b49977bcf3668d3870a1a", size = 9097190, upload-time = "2026-06-02T11:54:24.454Z" }, + { url = "https://files.pythonhosted.org/packages/fb/de/b650b4d69b84468cfa2e28a3ff7b8103743029e6446ce1a97fe060ef688c/scikit_learn-1.9.0-cp314-cp314t-win_amd64.whl", hash = "sha256:78fc56eafd4edb9575d2d8950d1dd152061abb573341a1cb7e099fc40f6c6666", size = 8963204, upload-time = "2026-06-02T11:54:27.428Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f3/ff83d76d7418112e5a61326443cdda87be3545dd8d6599c95b2481a4419e/scikit_learn-1.9.0-cp314-cp314t-win_arm64.whl", hash = "sha256:051075bda8b7aab87b1906ab3d4740a1e1224a19d7b3781a576736edc94e76aa", size = 8222661, upload-time = "2026-06-02T11:54:30.192Z" }, +] + [[package]] name = "scikit-network" version = "0.33.5" @@ -3127,6 +3434,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/07/39/338d9219c4e87f3e708f18857ecd24d22a0c3094752393319553096b98af/scipy-1.17.1-cp314-cp314t-win_arm64.whl", hash = "sha256:200e1050faffacc162be6a486a984a0497866ec54149a01270adc8a59b7c7d21", size = 25489165, upload-time = "2026-02-23T00:22:29.563Z" }, ] +[[package]] +name = "sentence-transformers" +version = "5.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "scikit-learn" }, + { name = "scipy" }, + { name = "torch" }, + { name = "tqdm" }, + { name = "transformers" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f9/56/d2cb00765a6b15c994a7fccf20f9032f16e8193ca49147cb5155166ad744/sentence_transformers-5.6.0.tar.gz", hash = "sha256:0e7164d051e416c1853ade7c274ff52af3f9da0f4be7f0b83d734c27699e1057", size = 453194, upload-time = "2026-06-16T14:01:56.42Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/76/c1/dc1582b79e9a2eb0cddf9559cd9bcdff084f541d6fe881fdd9d98630dba7/sentence_transformers-5.6.0-py3-none-any.whl", hash = "sha256:d2075b5e687a1611005e20ab04a6846994d51adfcf39610aed066af3c0c0b81f", size = 596411, upload-time = "2026-06-16T14:01:55.103Z" }, +] + +[[package]] +name = "setuptools" +version = "81.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, +] + [[package]] name = "shellingham" version = "1.5.4" @@ -3236,6 +3571,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/45/a132b9074aa18e799b891b91ad72133c98d8042c70f6240e4c5f9dabee2f/structlog-25.5.0-py3-none-any.whl", hash = "sha256:a8453e9b9e636ec59bd9e79bbd4a72f025981b3ba0f5837aebf48f02f37a7f9f", size = 72510, upload-time = "2025-10-27T08:28:21.535Z" }, ] +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + [[package]] name = "tenacity" version = "9.1.4" @@ -3245,6 +3592,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" }, ] +[[package]] +name = "threadpoolctl" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b7/4d/08c89e34946fce2aec4fbb45c9016efd5f4d7f24af8e5d93296e935631d8/threadpoolctl-3.6.0.tar.gz", hash = "sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e", size = 21274, upload-time = "2025-03-13T13:49:23.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" }, +] + [[package]] name = "tiktoken" version = "0.13.0" @@ -3287,29 +3643,64 @@ wheels = [ [[package]] name = "tokenizers" -version = "0.23.1" +version = "0.22.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c1/60/21f715d9faba5f5407ff759472ade058ec4a507ad62bcea47cb847239a73/tokenizers-0.23.1.tar.gz", hash = "sha256:1feeeadf865a7915adc25445dea30e9933e593c31bb96c277cee36de227c8bfa", size = 365748, upload-time = "2026-04-27T14:43:25.606Z" } +sdist = { url = "https://files.pythonhosted.org/packages/73/6f/f80cfef4a312e1fb34baf7d85c72d4411afde10978d4657f8cdd811d3ccc/tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917", size = 372115, upload-time = "2026-01-05T10:45:15.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/97/5dbfabf04c7e348e655e907ed27913e03db0923abb5dfdd120d7b25630e1/tokenizers-0.22.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c", size = 3100275, upload-time = "2026-01-05T10:41:02.158Z" }, + { url = "https://files.pythonhosted.org/packages/2e/47/174dca0502ef88b28f1c9e06b73ce33500eedfac7a7692108aec220464e7/tokenizers-0.22.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001", size = 2981472, upload-time = "2026-01-05T10:41:00.276Z" }, + { url = "https://files.pythonhosted.org/packages/d6/84/7990e799f1309a8b87af6b948f31edaa12a3ed22d11b352eaf4f4b2e5753/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7", size = 3290736, upload-time = "2026-01-05T10:40:32.165Z" }, + { url = "https://files.pythonhosted.org/packages/78/59/09d0d9ba94dcd5f4f1368d4858d24546b4bdc0231c2354aa31d6199f0399/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd", size = 3168835, upload-time = "2026-01-05T10:40:38.847Z" }, + { url = "https://files.pythonhosted.org/packages/47/50/b3ebb4243e7160bda8d34b731e54dd8ab8b133e50775872e7a434e524c28/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5", size = 3521673, upload-time = "2026-01-05T10:40:56.614Z" }, + { url = "https://files.pythonhosted.org/packages/e0/fa/89f4cb9e08df770b57adb96f8cbb7e22695a4cb6c2bd5f0c4f0ebcf33b66/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e", size = 3724818, upload-time = "2026-01-05T10:40:44.507Z" }, + { url = "https://files.pythonhosted.org/packages/64/04/ca2363f0bfbe3b3d36e95bf67e56a4c88c8e3362b658e616d1ac185d47f2/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b", size = 3379195, upload-time = "2026-01-05T10:40:51.139Z" }, + { url = "https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67", size = 3274982, upload-time = "2026-01-05T10:40:58.331Z" }, + { url = "https://files.pythonhosted.org/packages/1d/28/5f9f5a4cc211b69e89420980e483831bcc29dade307955cc9dc858a40f01/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4", size = 9478245, upload-time = "2026-01-05T10:41:04.053Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fb/66e2da4704d6aadebf8cb39f1d6d1957df667ab24cff2326b77cda0dcb85/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a", size = 9560069, upload-time = "2026-01-05T10:45:10.673Z" }, + { url = "https://files.pythonhosted.org/packages/16/04/fed398b05caa87ce9b1a1bb5166645e38196081b225059a6edaff6440fac/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a", size = 9899263, upload-time = "2026-01-05T10:45:12.559Z" }, + { url = "https://files.pythonhosted.org/packages/05/a1/d62dfe7376beaaf1394917e0f8e93ee5f67fea8fcf4107501db35996586b/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5", size = 10033429, upload-time = "2026-01-05T10:45:14.333Z" }, + { url = "https://files.pythonhosted.org/packages/fd/18/a545c4ea42af3df6effd7d13d250ba77a0a86fb20393143bbb9a92e434d4/tokenizers-0.22.2-cp39-abi3-win32.whl", hash = "sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92", size = 2502363, upload-time = "2026-01-05T10:45:20.593Z" }, + { url = "https://files.pythonhosted.org/packages/65/71/0670843133a43d43070abeb1949abfdef12a86d490bea9cd9e18e37c5ff7/tokenizers-0.22.2-cp39-abi3-win_amd64.whl", hash = "sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48", size = 2747786, upload-time = "2026-01-05T10:45:18.411Z" }, + { url = "https://files.pythonhosted.org/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, +] + +[[package]] +name = "torch" +version = "2.12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/87/39/b87a87d5bb9470610b80a2d31df42fcffeaf35118b8b97952b2aff598cc7/tokenizers-0.23.1-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:e03d6ffcbe0d56ee9c1ccd070e70a13fa750727c0277e138152acbc0252c2224", size = 3146732, upload-time = "2026-04-27T14:43:15.427Z" }, - { url = "https://files.pythonhosted.org/packages/e2/6a/068ed9f6e444c9d7e9d55ce134181325700f3d7f30410721bdc8f848d727/tokenizers-0.23.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:e0948bbb1ac1d7cdfc9fb6d62c596e3b7550036ad60ecd654a66ad273326324e", size = 3054954, upload-time = "2026-04-27T14:43:13.745Z" }, - { url = "https://files.pythonhosted.org/packages/6c/36/e006edf031154cba92b8416057d92c3abe3635e4c4b0aa0b5b9bb39dde70/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1bf13402aff9bc533c89cb849ec3b412dc3fbeacc9744840e423d7bf3f7dc0e3", size = 3374081, upload-time = "2026-04-27T14:43:01.241Z" }, - { url = "https://files.pythonhosted.org/packages/a2/ef/7735d226f9c7f874a6bee5e3f27fb25ecabdf207d37b8cf45286d0795893/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f836ca703b89ae07919a309f9651f7a88fd5a33d5f718ba5ad0870ec0256bad6", size = 3247641, upload-time = "2026-04-27T14:43:03.856Z" }, - { url = "https://files.pythonhosted.org/packages/b9/d9/24827036f6e21297bfffda0768e58eb6096a4f411e932964a01707857931/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ae848657742035523fdf261773630cb819a26995fcd3d9ecae0c1daf6e5a4959", size = 3585624, upload-time = "2026-04-27T14:43:10.664Z" }, - { url = "https://files.pythonhosted.org/packages/0c/9a/22f3582b3a4f49358293a5206e25317621ee4526bfe9cdaa0f07a12e770e/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:53b09e85775d5187941e7bab30e941b4134ab4a7dd8c68e783d231fb7ca27c51", size = 3844062, upload-time = "2026-04-27T14:43:05.643Z" }, - { url = "https://files.pythonhosted.org/packages/7e/65/b8f8814eef95800f20721384136d9a1d22241d50b2874357cb70542c392f/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ea5a0ce170074329faaa8ea3f6400ecde604b6678192688533af80980daae71a", size = 3460098, upload-time = "2026-04-27T14:43:08.854Z" }, - { url = "https://files.pythonhosted.org/packages/0d/d5/1353e5f677ec27c2494fb6a6725e82d56c985f53e90ec511369e7e4f02c6/tokenizers-0.23.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5075b405006415ea148a992d093699c66eb01952bf59f4d5727089a98bda45a4", size = 3346235, upload-time = "2026-04-27T14:43:12.377Z" }, - { url = "https://files.pythonhosted.org/packages/71/89/39b6b8fc073fb6d413d0147aa333dc7eff7be65639ac9d19930a0b21bf33/tokenizers-0.23.1-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:56f3a77de629917652f876294dc9fe6bad4a0c43bc229dc72e59bb23a0f4729a", size = 3426398, upload-time = "2026-04-27T14:43:07.264Z" }, - { url = "https://files.pythonhosted.org/packages/0f/80/127c854da64827e5b79264ce524993a90dddcb320e5cd42412c5c02f9e8a/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9d10a6d957ef01896dc274e890eee27d41bd0e74ef31e60616f0fc311345184e", size = 9823279, upload-time = "2026-04-27T14:43:17.222Z" }, - { url = "https://files.pythonhosted.org/packages/fe/ba/44c2502feb1a058f096ddfb4e0996ef3225a01a388e1a9b094e91689fe93/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:1974288a609c343774f1b897c8b482c791ab17b75ab5c8c2b1737565c1d82288", size = 9644986, upload-time = "2026-04-27T14:43:19.45Z" }, - { url = "https://files.pythonhosted.org/packages/9e/c1/464019a9fb059870bfe4eebb4ba12208f3042035e258bf5e782906bd3847/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:120468fb4c24faf0543c835a4fabafa4deb3f20a035c9b6e83d0b553a97615d4", size = 9976181, upload-time = "2026-04-27T14:43:21.463Z" }, - { url = "https://files.pythonhosted.org/packages/79/94/3ac1432bda31626071e9b6a12709b97ae05131c804b94c8f3ac622c5da32/tokenizers-0.23.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e3d8f40ea6268047de7046906326abed5134f27d4e8447b23763afe5808c8a96", size = 10113853, upload-time = "2026-04-27T14:43:23.617Z" }, - { url = "https://files.pythonhosted.org/packages/6a/dd/631b21433c771b1382535326f0eca80b9c9cee2e64961dd993bc9ac4669e/tokenizers-0.23.1-cp310-abi3-win32.whl", hash = "sha256:93120a930b919416da7cd10a2f606ac9919cc69cacae7980fa2140e277660948", size = 2536263, upload-time = "2026-04-27T14:43:29.888Z" }, - { url = "https://files.pythonhosted.org/packages/97/c9/2553f72aaf65a2797d4229e37fa7fbe38ffbf3e32912d31bdd78b3323e59/tokenizers-0.23.1-cp310-abi3-win_amd64.whl", hash = "sha256:e7bfaf995c1bdbbd21d13539decb6650967013759318627d85daeb7881af16b7", size = 2798223, upload-time = "2026-04-27T14:43:28.51Z" }, - { url = "https://files.pythonhosted.org/packages/cd/2b/2be299bab55fc595e3d38567edb1a87f86e594842968fa9515a07bdcf422/tokenizers-0.23.1-cp310-abi3-win_arm64.whl", hash = "sha256:a26197957d8e4425dfba746315f3c425ea00cfa8367c5fbc4ec73447893dcea9", size = 2664127, upload-time = "2026-04-27T14:43:26.949Z" }, + { url = "https://files.pythonhosted.org/packages/ff/4a/0300261818e1560d72cc160ac826005507e8b7ca0a35788b591436d05b4a/torch-2.12.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:c75e93173c700bccd6bfcc4a9d19ce242ab6dacd1f1781483027a16239b9e650", size = 87992358, upload-time = "2026-06-17T21:07:40.299Z" }, + { url = "https://files.pythonhosted.org/packages/30/a7/874a5ca05e8f159211dca7921060f7057acc1adb26431e119fd150623efc/torch-2.12.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:fcb61ccd20784b62bdd78ec84238a5cfb383b4994902e03bac95505ab360884c", size = 426386134, upload-time = "2026-06-17T21:07:31.481Z" }, + { url = "https://files.pythonhosted.org/packages/e1/75/20bb8fe9c1ad6538cce8cd0391b51927ae5af0b17ed1eab44b8824465dc1/torch-2.12.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f4afc8083dff08719edbea346644476e3cec0cf40ebe256be0ee5d5b7c7e8c0d", size = 532268019, upload-time = "2026-06-17T21:05:37.925Z" }, + { url = "https://files.pythonhosted.org/packages/d1/fa/824ddb662af55b2eabc0dbb7b57c7c0b1bcd93693754a2b8509ec4d16490/torch-2.12.1-cp313-cp313-win_amd64.whl", hash = "sha256:f92609e3b3ce72f25e2eb780d043ced2480c1a86c47c852604fc7a9108648386", size = 122987777, upload-time = "2026-06-17T21:07:09.49Z" }, + { url = "https://files.pythonhosted.org/packages/63/b7/1b49fe7086ea36839cc80abc43174c43d0ab6f676c0891c871c162f44fe3/torch-2.12.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:e9b6f7d2dd66ea87a3ae620069d31335d594c06effb1a383bdd21cfe61e44ece", size = 88010025, upload-time = "2026-06-17T21:07:03.934Z" }, + { url = "https://files.pythonhosted.org/packages/d7/06/5b44063a6545036dcc680d2d303b137d9176cfb2cc1e1863e3ef94abeb52/torch-2.12.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:7973ccd3d2cd35c74449213f7bded199bec6c6247e705cbeda7407af79703d91", size = 426392891, upload-time = "2026-06-17T21:05:52.261Z" }, + { url = "https://files.pythonhosted.org/packages/f8/dd/c9ce9a4b0eb3c5bb92d9ea56766e2c22559f0b45171149188494edcce80f/torch-2.12.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:c64ac4aac16be5e296dcd912305605804b203333c690bf98c55bc09494ee92ad", size = 532272494, upload-time = "2026-06-17T21:06:22.72Z" }, + { url = "https://files.pythonhosted.org/packages/21/7c/f3a601fc1b1f663ff269bfe553654e638651939aa6563e8daa7167c33098/torch-2.12.1-cp314-cp314-win_amd64.whl", hash = "sha256:f6dc4caf7eb4adb38a2d9f536b51db56310fdd1254e69a2d96767e1367c892b3", size = 122987254, upload-time = "2026-06-17T21:06:33.199Z" }, + { url = "https://files.pythonhosted.org/packages/e6/8c/b8087556cf81ddd808dbeb34afb8396d7ae7a1694ab489f08b1a0004e7d0/torch-2.12.1-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:2afbb2bdaa8a95040e733f05492ddf133c3967c9b7ce0abd218d704b6cab437d", size = 88303173, upload-time = "2026-06-17T21:05:06.603Z" }, + { url = "https://files.pythonhosted.org/packages/4a/07/fe09d1699fbed2afa10ebc692ff2b99d113f2605b6748cea633989e2789a/torch-2.12.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:97eba061fcb042fed191400b15568990073d67eaacaa6ee9b7ca01dd8b790fe9", size = 426404009, upload-time = "2026-06-17T21:04:57.557Z" }, + { url = "https://files.pythonhosted.org/packages/2e/f7/0ce4f6c1962c60ded7270e0a9eb560fb615c92b89d332cf9e3dff36d5ecc/torch-2.12.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:3867b861391701012adb2df93360efb88494dca245a185e3bb7624495cfe3f33", size = 532184292, upload-time = "2026-06-17T21:05:17.526Z" }, + { url = "https://files.pythonhosted.org/packages/70/db/e384c12aba30320ca92aaaf557456cbcb26f04b4df307728bb8f019f5000/torch-2.12.1-cp314-cp314t-win_amd64.whl", hash = "sha256:dd15595f8fc764cffde8c6361a3beb6ef69a028c851b1b3e70e077f615980d4e", size = 123231142, upload-time = "2026-06-17T21:05:27.061Z" }, ] [[package]] @@ -3324,6 +3715,39 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl", hash = "sha256:ee1e4c0e59148062281c49d80b25b67771a127c85fc9676d3be5f243206826bf", size = 78374, upload-time = "2026-02-03T17:35:50.982Z" }, ] +[[package]] +name = "transformers" +version = "5.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, + { name = "typer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ec/e1/720ff7ff666b04279fea5bb7ac3ef8675e98f0ddbc1b8cb8bc9f3889d62e/transformers-5.13.0.tar.gz", hash = "sha256:940c1428e42a4238f9ccf0cd41e63c590701aa63c19fd2ce3d7d602222d68495", size = 9195801, upload-time = "2026-07-03T16:05:39.362Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a5/3a/d99704c5effe10c6339c98cb236259161103e159bb99a78468b6729572ec/transformers-5.13.0-py3-none-any.whl", hash = "sha256:8adbc1d20bd5463cd6876b2eb7cb31971e1065788e7dc6bc12bab597a7c504b7", size = 11503730, upload-time = "2026-07-03T16:05:35.569Z" }, +] + +[[package]] +name = "triton" +version = "3.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/42/c5089d4d9327fcd1e862c599cc2927f39418f84dd11a84cb2ccff9d4787a/triton-3.7.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdbfc09d9ec58bc5e68321525653220de7515c199e7a8097a97c85e62b52cd0a", size = 184694629, upload-time = "2026-06-17T20:03:53.444Z" }, + { url = "https://files.pythonhosted.org/packages/07/42/2c3ac59253ae8892b6f307875263dd23dc875cdf732d3aea40d6d41fb7cb/triton-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58c0e131da05134a2a4788ccbcc0c1105cf0f54c8e98f19e34cd465396dc15eb", size = 197729241, upload-time = "2026-06-17T19:53:27.801Z" }, + { url = "https://files.pythonhosted.org/packages/40/71/e01aa7ad573883ed9456f130226babdec70b005e098c4d6226a6238e761b/triton-3.7.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe4ea396a06171f1f1f58cbd39c70b09294398f7dd7c620939bab54ad6f934fa", size = 184705764, upload-time = "2026-06-17T20:03:59.064Z" }, + { url = "https://files.pythonhosted.org/packages/a4/09/5683146fda6a2b569deb78ccfd8fbfea8bfe55f726b081c0a6bb18dd6f28/triton-3.7.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2020153b08280415ec0da6607834e79166442147e78e144df06b508c75b186d2", size = 197729537, upload-time = "2026-06-17T19:53:35.516Z" }, + { url = "https://files.pythonhosted.org/packages/e9/f8/448220c3092019f9fdfab39ec47985968181d67da34b44f6a7f6280a5cbb/triton-3.7.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c58e4c61f0c73b5dba3b5d19b4a7093c32f90dc18b2a7f121a7c16ccd31107b7", size = 184814760, upload-time = "2026-06-17T20:04:04.984Z" }, + { url = "https://files.pythonhosted.org/packages/f0/ac/229b7d4589d2e5937310e72c6d46e89599d16a4a12b479ffa1499fee8eb8/triton-3.7.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10ba85fa2cca4a2fbdeb36bf1cb082f2c252bda55bf9fccd74f65ec5bc647e68", size = 197824404, upload-time = "2026-06-17T19:53:42.772Z" }, +] + [[package]] name = "typer" version = "0.25.1" From 517fa9e9320f668c55c5b265ff397b93af26786a Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 06:46:30 +0200 Subject: [PATCH 08/20] feat: low-confidence refusal gate for retrieval scores below threshold Claude-Session: https://claude.ai/code/session_01Md71iujmpvewiisGX5u33H --- .env.example | 2 ++ docs/rag-retrieval.md | 19 +++++++++++++++++++ localrag/api/schemas.py | 8 ++++++++ localrag/api/service.py | 9 ++++++++- localrag/rag/engine.py | 22 +++++++++++++++++++++- localrag/settings.py | 6 +++++- tests/test_rag_engine.py | 34 ++++++++++++++++++++++++++++++++++ 7 files changed, 97 insertions(+), 3 deletions(-) diff --git a/.env.example b/.env.example index 1c13bb0..c508eee 100644 --- a/.env.example +++ b/.env.example @@ -27,6 +27,8 @@ OCR_LANGUAGE=eng OCR_MIN_CHARS_PER_PAGE=20 RAG_TOP_K=5 +# Below this top-context score, refuse instead of generating (0 disables). +RAG_MIN_CONTEXT_SCORE=0.0 RETRIEVAL_MODE=hybrid BM25_WEIGHT=0.5 RRF_K=60 diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index 770d5cf..57a5438 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -124,6 +124,25 @@ builds the `CrossEncoderReranker` only when `RERANK_ENABLED=true`, mirroring the pluggable-provider shape in `localrag/llm/factory.py` (nothing imports `sentence-transformers` unless the feature is turned on). +## Low-confidence refusal gate + +`RAGEngine.stream_chat_from_contexts` (`localrag/rag/engine.py`) can short-circuit +before calling the LLM at all: if the top retrieved context's `score` is below +`RAG_MIN_CONTEXT_SCORE`, or no contexts were retrieved while the gate is enabled, +`_is_low_confidence` returns `True` and `_low_confidence_response` yields a single +canned refusal token plus a `final` event with `sources: []` and +`low_confidence: True`. Otherwise generation proceeds as usual and the `final` +event carries `low_confidence: False`. The flag is surfaced end-to-end: +`QueryResponse.low_confidence` in `POST /query`, and in the `final` SSE payload +for `POST /query/stream`. + +Disabled by default (`RAG_MIN_CONTEXT_SCORE=0.0`). Because the score scale +depends on the embedding model and `RETRIEVAL_MODE` (raw cosine/L2 similarity +vs. fused hybrid/RRF scores), there is no universal threshold — tune it +per-corpus after inspecting typical top scores for known-good and known-bad +queries. This is a lightweight heuristic gate, not a replacement for a real +guardrails layer (e.g. NeMo Guardrails) in a regulated setting. + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: diff --git a/localrag/api/schemas.py b/localrag/api/schemas.py index d514e7d..24f273c 100644 --- a/localrag/api/schemas.py +++ b/localrag/api/schemas.py @@ -247,6 +247,14 @@ class QueryResponse(BaseModel): description="Total wall-clock latency in milliseconds.", examples=[312.5] ) model: str = Field(description="Model tag used to generate the answer.", examples=["llama3.2"]) + low_confidence: bool = Field( + default=False, + description=( + "True when the top retrieved context scored below `RAG_MIN_CONTEXT_SCORE` and " + "the answer is a canned refusal instead of a generated response." + ), + examples=[False], + ) class AgentQueryRequest(BaseModel): diff --git a/localrag/api/service.py b/localrag/api/service.py index 387779d..ce01acd 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -286,6 +286,7 @@ def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: raise RagApiError(int(exc.status_code), exc.detail) from exc answer_chunks: list[str] = [] + low_confidence = False for event in engine.stream_chat_from_contexts( contexts=contexts, question=request.question, @@ -293,6 +294,8 @@ def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: ): if event["type"] == "token": answer_chunks.append(str(event["token"])) + if event["type"] == "final": + low_confidence = bool(event.get("low_confidence", False)) latency_ms = (time.perf_counter() - t0) * 1000 used_model = request.model or engine.settings.ollama_llm_model @@ -317,6 +320,7 @@ def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: sources=sources, latency_ms=latency_ms, model=used_model, + low_confidence=low_confidence, ) @@ -352,5 +356,8 @@ def iter_query_sse_events( if event["type"] == "token": yield {"event": "token", "data": str(event["token"])} if event["type"] == "final": - payload = {"sources": event["sources"]} + payload = { + "sources": event["sources"], + "low_confidence": event.get("low_confidence", False), + } yield {"event": "final", "data": json.dumps(payload)} diff --git a/localrag/rag/engine.py b/localrag/rag/engine.py index 0808c02..13de83e 100644 --- a/localrag/rag/engine.py +++ b/localrag/rag/engine.py @@ -63,8 +63,28 @@ def stream_chat_from_contexts( model: str | None, ) -> Generator[dict[str, Any]]: """Stream LLM tokens when contexts were retrieved earlier (HTTP runs retrieve first).""" + if self._is_low_confidence(contexts): + logger.info("rag_low_confidence_refusal question_chars=%s", len(question)) + return self._low_confidence_response() return self._stream_chat_tokens(contexts=contexts, question=question, model=model) + def _is_low_confidence(self, contexts: list[dict[str, Any]]) -> bool: + min_score = self.settings.rag_min_context_score + if min_score <= 0: + return False + if not contexts: + return True + top_score = float(contexts[0].get("score", 0.0)) + return top_score < min_score + + @staticmethod + def _low_confidence_response() -> Generator[dict[str, Any]]: + yield { + "type": "token", + "token": "I don't have enough information in the ingested documents to answer that.", + } + yield {"type": "final", "sources": [], "low_confidence": True} + def _stream_chat_tokens( self, *, @@ -82,7 +102,7 @@ def _stream_chat_tokens( if event["type"] == "token": yield event logger.info("rag_stream_done") - yield {"type": "final", "sources": self._extract_sources(contexts)} + yield {"type": "final", "sources": self._extract_sources(contexts), "low_confidence": False} @staticmethod def _extract_sources(contexts: list[dict[str, Any]]) -> list[dict[str, object]]: diff --git a/localrag/settings.py b/localrag/settings.py index 21152e8..970a28e 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -47,7 +47,10 @@ class Settings(BaseSettings): ``rag_system_prompt`` is the system message for the answering model. When ``parent_expansion_enabled`` is true (default), top hits with a non-empty ``heading_path`` are expanded to their full sibling-chunk - section before prompting; set false to disable. + section before prompting; set false to disable. ``rag_min_context_score`` + gates generation on retrieval confidence: below this score (or with no + contexts at all) the engine returns a canned refusal instead of calling + the LLM; ``0`` (default) disables the gate. **Reranking** — When ``rerank_enabled`` is true (default false, requires ``uv sync --extra rerank``), retrieval over-fetches ``rerank_fetch_k`` @@ -93,6 +96,7 @@ class Settings(BaseSettings): ocr_min_chars_per_page: int = 20 rag_top_k: int = 5 + rag_min_context_score: float = 0.0 retrieval_mode: str = "hybrid" bm25_weight: float = 0.5 rrf_k: int = 60 diff --git a/tests/test_rag_engine.py b/tests/test_rag_engine.py index 82e660f..3bce91d 100644 --- a/tests/test_rag_engine.py +++ b/tests/test_rag_engine.py @@ -142,3 +142,37 @@ def count_tokens(self, text: str) -> int: with pytest.raises(RuntimeError, match="provider down"): list(engine.stream_answer(question="Q", model="llm", n_results=1)) + + +def test_rag_engine_short_circuits_on_low_confidence_context() -> None: + settings = Settings(rag_min_context_score=0.5) + contexts = [{"text": "weak match", "source": "a.md", "chunk_index": 0, "score": 0.1}] + provider = FakeProvider(tokens=["should not be used"]) + engine = RAGEngine( + settings=settings, retriever=StubRetriever(contexts=contexts), provider=provider + ) + + events = list(engine.stream_answer(question="Q", n_results=1)) + + token_events = [ev for ev in events if ev["type"] == "token"] + expected_refusal = "I don't have enough information in the ingested documents to answer that." + assert token_events[0]["token"] == expected_refusal + final = events[-1] + assert final["low_confidence"] is True + assert final["sources"] == [] + assert provider.prompts_seen == [] + + +def test_rag_engine_empty_contexts_are_low_confidence_when_threshold_enabled() -> None: + settings = Settings(rag_min_context_score=0.1) + provider = FakeProvider(tokens=[]) + engine = RAGEngine(settings=settings, retriever=StubRetriever(contexts=[]), provider=provider) + + events = list(engine.stream_answer(question="Q", n_results=1)) + final = events[-1] + assert final["low_confidence"] is True + + +def test_rag_engine_low_confidence_disabled_by_default() -> None: + settings = Settings() + assert settings.rag_min_context_score == 0.0 From 87e54c1a6408f6cb3b12ec5b11575e9a37512c12 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 06:52:33 +0200 Subject: [PATCH 09/20] feat: optional LLM query rewriting for retrieval --- .env.example | 2 + docs/rag-retrieval.md | 23 ++++++++++ localrag/rag/query_rewrite.py | 35 +++++++++++++++ localrag/rag/retriever.py | 17 +++++-- localrag/settings.py | 6 +++ tests/test_query_rewrite.py | 85 +++++++++++++++++++++++++++++++++++ tests/test_retriever.py | 26 +++++++++++ 7 files changed, 191 insertions(+), 3 deletions(-) create mode 100644 localrag/rag/query_rewrite.py create mode 100644 tests/test_query_rewrite.py diff --git a/.env.example b/.env.example index c508eee..657af6b 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,8 @@ RRF_K=60 FRESHNESS_HALF_LIFE_DAYS=30.0 # Expand top retrieval hits to their full markdown section (via heading_path) before prompting. Set false to disable. PARENT_EXPANSION_ENABLED=true +# LLM rewrites the query for retrieval only (adds one extra LLM round-trip per query); the original question is still used for the final answer prompt. +QUERY_REWRITE_ENABLED=false RAG_SYSTEM_PROMPT=You are a helpful assistant. Answer only based on the provided context. # Optional cross-encoder reranking (requires `uv sync --extra rerank`). diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index 57a5438..33d24c1 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -143,6 +143,29 @@ per-corpus after inspecting typical top scores for known-good and known-bad queries. This is a lightweight heuristic gate, not a replacement for a real guardrails layer (e.g. NeMo Guardrails) in a regulated setting. +## Query rewriting (optional) + +Disabled by default (`QUERY_REWRITE_ENABLED=false`). When enabled, +`Retriever.retrieve` (`localrag/rag/retriever.py`) calls +`rewrite_query(question, settings)` (`localrag/rag/query_rewrite.py`) before +embedding/BM25 search. `rewrite_query` reuses +`localrag/llm/factory.py::build_provider` (the same `ResilientProvider`-wrapped +provider abstraction used for answering) with a retrieval-specific system +prompt that asks for a short, keyword-dense reformulation of the question, +preserving exact identifiers/codes/names verbatim. + +The rewrite is retrieval-only: it replaces the text sent to +`OllamaEmbedder.embed_text` and `Bm25Index.query`, but the original question is +still what gets passed to the reranker (if enabled) and to the final answer +prompt — rewriting never affects citations or the text the LLM sees when +generating the answer. On any provider failure (timeout, exception, empty +response) `rewrite_query` logs and falls back to the original question, so +retrieval degrades to its normal behavior rather than failing the request. + +This adds one extra LLM round-trip per query, so it is off by default; enable +it when lexical/embedding mismatch between conversational questions and +indexed document phrasing is hurting retrieval recall. + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: diff --git a/localrag/rag/query_rewrite.py b/localrag/rag/query_rewrite.py new file mode 100644 index 0000000..3dedc13 --- /dev/null +++ b/localrag/rag/query_rewrite.py @@ -0,0 +1,35 @@ +"""Optional LLM-based query rewriting for retrieval (not for the final answer prompt).""" + +from __future__ import annotations + +import logging + +from localrag.llm.factory import build_provider +from localrag.settings import Settings + +logger = logging.getLogger(__name__) + +_REWRITE_INSTRUCTION = ( + "Rewrite the user's question as a short, keyword-dense search query for a " + "document retrieval system. Keep any exact identifiers, codes, or names " + "verbatim. Respond with only the rewritten query, no explanation." +) + + +def rewrite_query(question: str, settings: Settings) -> str: + """Return a keyword-dense reformulation of ``question`` for retrieval only. + + Falls back to the original ``question`` on any provider failure. The + original question — not this rewrite — is still used for the final + answer prompt; rewriting is a retrieval-only concern. + """ + rewrite_provider = build_provider( + settings.model_copy(update={"rag_system_prompt": _REWRITE_INSTRUCTION}) + ) + try: + response = rewrite_provider.generate(question, context=[]) + except Exception: + logger.exception("query_rewrite_failed_falling_back_to_original") + return question + rewritten = response.answer.strip() + return rewritten or question diff --git a/localrag/rag/retriever.py b/localrag/rag/retriever.py index 234cf9d..478e7e4 100644 --- a/localrag/rag/retriever.py +++ b/localrag/rag/retriever.py @@ -11,6 +11,7 @@ from localrag.ingestion.embedder import OllamaEmbedder from localrag.rag.bm25_index import Bm25Index from localrag.rag.exceptions import RetrievalError +from localrag.rag.query_rewrite import rewrite_query from localrag.rag.reranker import CrossEncoderReranker from localrag.settings import Settings from localrag.storage.vector_store import VectorStore @@ -44,9 +45,19 @@ def retrieve( if self.reranker is not None else max(top_k * 2, top_k) ) - logger.debug("retrieve_embed_question top_k=%s question_chars=%s", top_k, len(question)) + search_question = question + if self.settings.query_rewrite_enabled: + search_question = rewrite_query(question, self.settings) + logger.debug( + "query_rewritten original_chars=%s rewritten_chars=%s", + len(question), + len(search_question), + ) + logger.debug( + "retrieve_embed_question top_k=%s question_chars=%s", top_k, len(search_question) + ) try: - embedding = self.embedder.embed_text(question) + embedding = self.embedder.embed_text(search_question) except httpx.HTTPError as exc: logger.error( "retrieve_embed_ollama_http_error url=%s error=%s", @@ -76,7 +87,7 @@ def retrieve( "ingested_at": hit.metadata.get("ingested_at"), "metadata": hit.metadata, } - for hit in self.bm25_index.query(question, top_k=fetch_k) + for hit in self.bm25_index.query(search_question, top_k=fetch_k) if _matches_filter(hit.metadata, metadata_filter) ] candidates = self._fuse_results( diff --git a/localrag/settings.py b/localrag/settings.py index 970a28e..fe7f85f 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -57,6 +57,11 @@ class Settings(BaseSettings): candidates and a local cross-encoder (``rerank_model``) re-scores and trims them to ``rag_top_k`` before freshness/expansion. + **Query rewriting** — When ``query_rewrite_enabled`` is true (default + false), an extra LLM round-trip rewrites the question into a keyword-dense + search query before embedding/BM25 retrieval; the original question is + still used for the final answer prompt. + **API** — ``api_host`` / ``api_port`` are the uvicorn bind address and port. **Logging** — ``log_level`` is the minimum level for the ``localrag`` logger @@ -102,6 +107,7 @@ class Settings(BaseSettings): rrf_k: int = 60 freshness_half_life_days: float = 30.0 parent_expansion_enabled: bool = True + query_rewrite_enabled: bool = False rag_system_prompt: str = ( "You are a helpful assistant. Answer only based on the provided context." ) diff --git a/tests/test_query_rewrite.py b/tests/test_query_rewrite.py new file mode 100644 index 0000000..89433ed --- /dev/null +++ b/tests/test_query_rewrite.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from collections.abc import Generator +from dataclasses import dataclass +from typing import Any + +import pytest + +from localrag.llm.providers.base import BaseLLMProvider +from localrag.llm.types import LLMResponse +from localrag.rag.query_rewrite import rewrite_query +from localrag.settings import Settings + + +@dataclass +class FakeProvider(BaseLLMProvider): + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + assert prompt == "how do i fix ERR_QUIC_PROTOCOL_ERROR" + return LLMResponse( + answer="ERR_QUIC_PROTOCOL_ERROR fix", + model="m", + tokens_used=3, + latency_ms=1.0, + estimated_cost_usd=0.0, + ) + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError + + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + raise NotImplementedError + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError + + def count_tokens(self, text: str) -> int: + raise NotImplementedError + + +@dataclass +class ExplodingProvider(BaseLLMProvider): + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + raise RuntimeError("boom") + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError + + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + raise NotImplementedError + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError + + def count_tokens(self, text: str) -> int: + raise NotImplementedError + + +def test_rewrite_query_returns_provider_answer(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "localrag.rag.query_rewrite.build_provider", lambda _settings: FakeProvider() + ) + + out = rewrite_query("how do i fix ERR_QUIC_PROTOCOL_ERROR", Settings()) + + assert out == "ERR_QUIC_PROTOCOL_ERROR fix" + + +def test_rewrite_query_falls_back_to_original_on_provider_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "localrag.rag.query_rewrite.build_provider", lambda _settings: ExplodingProvider() + ) + + out = rewrite_query("original question", Settings()) + + assert out == "original question" diff --git a/tests/test_retriever.py b/tests/test_retriever.py index bb43393..44ecac3 100644 --- a/tests/test_retriever.py +++ b/tests/test_retriever.py @@ -217,3 +217,29 @@ def query( contexts = retriever.retrieve("q") assert "expanded_text" not in contexts[0] + + +def test_retriever_uses_rewritten_query_for_embedding_when_enabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "localrag.rag.retriever.rewrite_query", + lambda question, _settings: "rewritten " + question, + ) + seen_questions: list[str] = [] + + @dataclass + class RecordingEmbedder: + def embed_text(self, text: str, *, model: str | None = None) -> list[float]: + seen_questions.append(text) + return [0.1, 0.2, 0.3] + + retriever = Retriever( + settings=Settings(query_rewrite_enabled=True), + embedder=RecordingEmbedder(), # type: ignore[arg-type] + vector_store=StubStore(), # type: ignore[arg-type] + ) + + retriever.retrieve("original question") + + assert seen_questions == ["rewritten original question"] From 3693ef6056b2031378f5c288193873f491edb5cf Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 06:53:11 +0200 Subject: [PATCH 10/20] feat: enrich SourceRef with heading_path and chunk_type for traceability Claude-Session: https://claude.ai/code/session_01Md71iujmpvewiisGX5u33H --- localrag/api/schemas.py | 13 +++++++++++++ localrag/api/service.py | 7 ++++++- localrag/rag/engine.py | 10 +++++++++- tests/test_rag_engine.py | 32 +++++++++++++++++++++++++++++++- 4 files changed, 59 insertions(+), 3 deletions(-) diff --git a/localrag/api/schemas.py b/localrag/api/schemas.py index 24f273c..cf146d1 100644 --- a/localrag/api/schemas.py +++ b/localrag/api/schemas.py @@ -220,6 +220,19 @@ class SourceRef(BaseModel): chunk_index: int = Field( description="Zero-based index of the chunk within the source.", examples=[0] ) + heading_path: str | None = Field( + default=None, + description="Markdown heading breadcrumb for this chunk, e.g. 'Setup > Install'.", + examples=["Setup > Install"], + ) + chunk_type: str | None = Field( + default=None, + description=( + "Chunk classification from ingestion, e.g. markdown_section, markdown_table, " + "code_block." + ), + examples=["markdown_section"], + ) class QueryResponse(BaseModel): diff --git a/localrag/api/service.py b/localrag/api/service.py index ce01acd..b96c5ee 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -301,7 +301,12 @@ def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: used_model = request.model or engine.settings.ollama_llm_model raw_sources = engine._extract_sources(contexts) # noqa: SLF001 sources = [ - SourceRef(source=str(s.get("source", "")), chunk_index=int(s.get("chunk_index", -1))) # type: ignore[call-overload] + 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 ] diff --git a/localrag/rag/engine.py b/localrag/rag/engine.py index 13de83e..552c2bc 100644 --- a/localrag/rag/engine.py +++ b/localrag/rag/engine.py @@ -115,5 +115,13 @@ def _extract_sources(contexts: list[dict[str, Any]]) -> list[dict[str, object]]: if key in seen: continue seen.add(key) - sources.append({"source": source, "chunk_index": chunk_index}) + metadata = context.get("metadata") or {} + sources.append( + { + "source": source, + "chunk_index": chunk_index, + "heading_path": metadata.get("heading_path") or None, + "chunk_type": metadata.get("chunk_type") or None, + } + ) return sources diff --git a/tests/test_rag_engine.py b/tests/test_rag_engine.py index 3bce91d..714b58f 100644 --- a/tests/test_rag_engine.py +++ b/tests/test_rag_engine.py @@ -73,11 +73,41 @@ def test_rag_engine_stream_answer_yields_tokens_and_dedupes_sources() -> None: final = events[-1] assert final["type"] == "final" - assert final["sources"] == [{"source": "a.md", "chunk_index": 1}] + assert final["sources"] == [ + {"source": "a.md", "chunk_index": 1, "heading_path": None, "chunk_type": None} + ] assert "SYS" in provider.prompts_seen[0] assert "chunk-one" in provider.prompts_seen[0] +def test_rag_engine_extract_sources_includes_heading_path_and_chunk_type() -> None: + settings = Settings(rag_system_prompt="SYS") + contexts = [ + { + "text": "chunk", + "source": "guide.md", + "chunk_index": 2, + "metadata": {"heading_path": "Setup > Install", "chunk_type": "markdown_section"}, + } + ] + engine = RAGEngine( + settings=settings, + retriever=StubRetriever(contexts=contexts), + provider=FakeProvider(tokens=[]), + ) + + sources = engine._extract_sources(contexts) # noqa: SLF001 + + assert sources == [ + { + "source": "guide.md", + "chunk_index": 2, + "heading_path": "Setup > Install", + "chunk_type": "markdown_section", + } + ] + + def test_rag_engine_answer_concatenates_tokens_and_returns_sources( monkeypatch: pytest.MonkeyPatch, ) -> None: From 4c749a24d0268bbf453cd6364e2e0bdd91c67469 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 06:56:46 +0200 Subject: [PATCH 11/20] feat: durable local audit log for queries (opt-in via AUDIT_LOG_PATH) --- .env.example | 3 +++ docs/architecture.md | 1 + localrag/api/service.py | 29 ++++++++++++++++++++++-- localrag/audit.py | 42 +++++++++++++++++++++++++++++++++++ localrag/settings.py | 2 ++ tests/test_audit.py | 49 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 localrag/audit.py create mode 100644 tests/test_audit.py diff --git a/.env.example b/.env.example index 657af6b..4f603e4 100644 --- a/.env.example +++ b/.env.example @@ -20,6 +20,9 @@ INGEST_ROOTS=[] UPLOAD_DIR=./data/uploads UPLOAD_MAX_BYTES=100000000 +# Append-only local audit trail (one JSON line per query: question, sources, answer, model, latency). Empty disables. +AUDIT_LOG_PATH= + # OCR fallback for scanned/image-only PDF pages (requires the `tesseract` binary # on the host; see docs/ocr.md). Set OCR_ENABLED=false to disable. OCR_ENABLED=true diff --git a/docs/architecture.md b/docs/architecture.md index a41e614..2730709 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -71,6 +71,7 @@ flowchart LR | LLM abstraction | `localrag/llm/` | `BaseLLMProvider`, Ollama/OpenAI/Anthropic providers, factory, cost estimator | | Agent | `localrag/agent/service.py`, `localrag/api/routers/agent.py` | Anthropic tool-use agent; `POST /agent/query` | | Eval | `evals/dataset.json`, `evals/run_evals.py` | RAGAS evaluation dataset and runner; `localrag eval` CLI command | +| Audit log | `localrag/audit.py` | `write_audit_record` — durable local JSONL trail (question, sources, answer, model, latency); disabled by default via `AUDIT_LOG_PATH` | ## LLM abstraction diff --git a/localrag/api/service.py b/localrag/api/service.py index b96c5ee..537301f 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -30,8 +30,10 @@ RebuildCollectionResponse, SourceRef, ) +from localrag.audit import write_audit_record from localrag.ingestion.loader import SUPPORTED_EXTENSIONS from localrag.ingestion.service import IngestionResult, IngestionService +from localrag.logging_config import request_id_ctx from localrag.ollama.schemas import OllamaTagsResponse, parse_ollama_json from localrag.rag.engine import RAGEngine from localrag.rag.exceptions import RetrievalError @@ -320,13 +322,23 @@ def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: latency_ms, len(sources), ) - return QueryResponse( + response = QueryResponse( answer="".join(answer_chunks).strip(), sources=sources, latency_ms=latency_ms, model=used_model, low_confidence=low_confidence, ) + write_audit_record( + engine.settings.audit_log_path, + correlation_id=request_id_ctx.get(""), + question=request.question, + sources=[s.model_dump() for s in sources], + answer=response.answer, + model=used_model, + latency_ms=latency_ms, + ) + return response def get_query_contexts(request: QueryRequest, engine: RAGEngine) -> list[dict[str, Any]]: @@ -346,6 +358,7 @@ def iter_query_sse_events( engine: RAGEngine, contexts: list[dict[str, Any]], ) -> Iterator[dict[str, Any]]: + t0 = time.perf_counter() logger.info( "query_start model=%s n_results=%s question_chars=%s", request.model, @@ -357,12 +370,24 @@ def iter_query_sse_events( question=request.question, model=request.model, ) + answer_chunks: list[str] = [] for event in stream: if event["type"] == "token": - yield {"event": "token", "data": str(event["token"])} + token = str(event["token"]) + answer_chunks.append(token) + yield {"event": "token", "data": token} if event["type"] == "final": payload = { "sources": event["sources"], "low_confidence": event.get("low_confidence", False), } + write_audit_record( + engine.settings.audit_log_path, + correlation_id=request_id_ctx.get(""), + question=request.question, + sources=event["sources"], + answer="".join(answer_chunks).strip(), + model=request.model or engine.settings.ollama_llm_model, + latency_ms=(time.perf_counter() - t0) * 1000, + ) yield {"event": "final", "data": json.dumps(payload)} diff --git a/localrag/audit.py b/localrag/audit.py new file mode 100644 index 0000000..4b6d908 --- /dev/null +++ b/localrag/audit.py @@ -0,0 +1,42 @@ +"""Durable local audit trail for RAG queries (not a regulatory-compliance system).""" + +from __future__ import annotations + +import json +import logging +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + + +def write_audit_record( + audit_log_path: str, + *, + correlation_id: str, + question: str, + sources: list[dict[str, Any]], + answer: str, + model: str, + latency_ms: float, +) -> None: + """Append one JSON line to ``audit_log_path``. No-op when the path is empty.""" + if not audit_log_path: + return + record = { + "timestamp": datetime.now(UTC).isoformat(), + "correlation_id": correlation_id, + "question": question, + "sources": sources, + "answer": answer, + "model": model, + "latency_ms": latency_ms, + } + try: + path = Path(audit_log_path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record) + "\n") + except OSError: + logger.exception("audit_log_write_failed path=%s", audit_log_path) diff --git a/localrag/settings.py b/localrag/settings.py index fe7f85f..fa1d19b 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -96,6 +96,8 @@ class Settings(BaseSettings): upload_dir: str = "./data/uploads" upload_max_bytes: int = 100_000_000 + audit_log_path: str = "" + ocr_enabled: bool = True ocr_language: str = "eng" ocr_min_chars_per_page: int = 20 diff --git a/tests/test_audit.py b/tests/test_audit.py new file mode 100644 index 0000000..41df4c2 --- /dev/null +++ b/tests/test_audit.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from localrag.audit import write_audit_record + + +def test_write_audit_record_appends_json_line(tmp_path: Path) -> None: + log_path = tmp_path / "audit.jsonl" + + write_audit_record( + str(log_path), + correlation_id="rid-1", + question="What is X?", + sources=[{"source": "a.md", "chunk_index": 0}], + answer="X is Y.", + model="llama3.2", + latency_ms=123.4, + ) + write_audit_record( + str(log_path), + correlation_id="rid-2", + question="Second question", + sources=[], + answer="Second answer", + model="llama3.2", + latency_ms=50.0, + ) + + lines = log_path.read_text(encoding="utf-8").strip().splitlines() + assert len(lines) == 2 + first = json.loads(lines[0]) + assert first["correlation_id"] == "rid-1" + assert first["question"] == "What is X?" + assert first["answer"] == "X is Y." + + +def test_write_audit_record_noop_when_path_empty(tmp_path: Path) -> None: + write_audit_record( + "", + correlation_id="rid", + question="q", + sources=[], + answer="a", + model="m", + latency_ms=1.0, + ) + assert list(tmp_path.iterdir()) == [] From 20eeb94bf95586a80a254ad6532f4ca68cb699fd Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 06:57:10 +0200 Subject: [PATCH 12/20] feat: optional tenant_id metadata tag for shared-deployment query scoping Claude-Session: https://claude.ai/code/session_01Md71iujmpvewiisGX5u33H --- .env.example | 7 +++++ docs/rag-retrieval.md | 22 +++++++++++++++ localrag/ingestion/service.py | 1 + localrag/settings.py | 10 +++++++ tests/test_ingestion_service.py | 47 +++++++++++++++++++++++++++++++++ 5 files changed, 87 insertions(+) diff --git a/.env.example b/.env.example index 4f603e4..e7bafc3 100644 --- a/.env.example +++ b/.env.example @@ -5,6 +5,13 @@ OLLAMA_LLM_MODEL=gemma3:4b CHROMA_PERSIST_PATH=./data/chroma CHROMA_COLLECTION_NAME=localrag +# Optional tenant tag written to every chunk's metadata and usable as a +# QueryRequest.metadata_filter key ({"tenant_id": "..."}). Empty = untagged +# (single-tenant deployments, the common case, pay zero extra cost). NOTE: +# this is an equality-filter convenience, not full multi-tenant isolation — +# see docs/rag-retrieval.md. +TENANT_ID= + CHUNK_CHARS=512 CHUNK_OVERLAP_CHARS=150 CHUNKING_MODE=structural diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index 33d24c1..a5138e2 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -166,6 +166,27 @@ This adds one extra LLM round-trip per query, so it is off by default; enable it when lexical/embedding mismatch between conversational questions and indexed document phrasing is hurting retrieval recall. +## Tenant tagging (optional) + +Per Chroma's own multi-tenancy guidance, this project uses a `tenant_id` +metadata field filtered at query time (via the metadata pre-filtering above) +rather than per-tenant Chroma collections, which the Chroma Cookbook +explicitly warns fragments the HNSW index and breaks whole-collection +operations like `Bm25Index.from_vector_store`. + +`TENANT_ID` (`localrag/settings.py`, empty by default) is written to every +chunk's metadata at ingest time (`localrag/ingestion/service.py::_ingest_one`). +A caller scopes retrieval to one tenant by passing +`{"question": "...", "metadata_filter": {"tenant_id": "team-a"}}` to +`POST /query` — no new retrieval code is needed since this reuses the +`metadata_filter` mechanism described above. + +This is an equality-filter convenience for a small-team shared deployment, not +a security boundary — anyone with API access can still query across all +`tenant_id` values by omitting the filter; pair with `API_KEY` and, if genuine +per-tenant isolation is ever required, revisit as a dedicated +(out-of-scope-for-this-plan) access-control project. + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: @@ -178,6 +199,7 @@ Freshness and debugging depend on chunk metadata written during ingestion: - `content_hash` - `source_mtime` - `git_commit` +- `tenant_id` The retriever returns `freshness_factor` and `ingested_at` in contexts so rank decisions are visible in API and test traces. diff --git a/localrag/ingestion/service.py b/localrag/ingestion/service.py index d43a5dc..38ca50e 100644 --- a/localrag/ingestion/service.py +++ b/localrag/ingestion/service.py @@ -203,6 +203,7 @@ def _ingest_one(self, resolved_path: Path, embed_model: str | None) -> int | Non "content_hash": content_hash, "source_mtime": source_mtime, "git_commit": git_commit, + "tenant_id": self.settings.tenant_id, } for index, chunk in enumerate(structural_chunks) ] diff --git a/localrag/settings.py b/localrag/settings.py index fa1d19b..1497335 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -67,6 +67,11 @@ class Settings(BaseSettings): **Logging** — ``log_level`` is the minimum level for the ``localrag`` logger (``DEBUG``, ``INFO``, ``WARNING``, ``ERROR``). Used when the API starts and when the CLI process starts. + + **Tenant tagging** — ``tenant_id`` (empty by default) is written to every + chunk's metadata at ingest time and can be used as a + ``QueryRequest.metadata_filter`` key (``{"tenant_id": "..."}``) to scope + retrieval to one tenant. See `docs/rag-retrieval.md`. """ model_config = SettingsConfigDict( @@ -152,6 +157,11 @@ class Settings(BaseSettings): log_level: str = "INFO" + # Optional tenant tag written to every chunk's metadata and usable as a + # QueryRequest.metadata_filter key ({"tenant_id": "..."}). Empty = untagged + # (single-tenant deployments, the common case, pay zero extra cost). + tenant_id: str = "" + @lru_cache(maxsize=1) def get_settings() -> Settings: diff --git a/tests/test_ingestion_service.py b/tests/test_ingestion_service.py index 34d9ec3..dcbfa34 100644 --- a/tests/test_ingestion_service.py +++ b/tests/test_ingestion_service.py @@ -333,6 +333,53 @@ def test_ingestion_service_ingest_one_writes_content_hash_and_git_metadata(tmp_p assert all(md["git_commit"] == "" for md in metadatas) # type: ignore[union-attr] +def test_ingestion_service_writes_tenant_id_from_settings(tmp_path: Path) -> None: + allowed_root = tmp_path / "allowed" + allowed_root.mkdir() + path = allowed_root / "a.md" + path.write_text("hello world", encoding="utf-8") + + settings = Settings( + ingest_roots=[str(allowed_root)], + chunk_chars=100, + chunk_overlap_chars=0, + tenant_id="team-a", + ) + embedder = StubEmbedder(seen_texts_batches=[]) + vector_store = StubVectorStore(deleted_sources=[], added=[], distinct_sources=None) + service = IngestionService( + settings=settings, + embedder=embedder, # type: ignore[arg-type] + vector_store=vector_store, # type: ignore[arg-type] + ) + + service.ingest_paths([path]) + + metadatas = vector_store.added[0]["metadatas"] + assert all(md["tenant_id"] == "team-a" for md in metadatas) # type: ignore[union-attr] + + +def test_ingestion_service_tenant_id_defaults_to_empty_string(tmp_path: Path) -> None: + allowed_root = tmp_path / "allowed" + allowed_root.mkdir() + path = allowed_root / "a.md" + path.write_text("hello world", encoding="utf-8") + + settings = Settings(ingest_roots=[str(allowed_root)], chunk_chars=100, chunk_overlap_chars=0) + embedder = StubEmbedder(seen_texts_batches=[]) + vector_store = StubVectorStore(deleted_sources=[], added=[], distinct_sources=None) + service = IngestionService( + settings=settings, + embedder=embedder, # type: ignore[arg-type] + vector_store=vector_store, # type: ignore[arg-type] + ) + + service.ingest_paths([path]) + + metadatas = vector_store.added[0]["metadatas"] + assert all(md["tenant_id"] == "" for md in metadatas) # type: ignore[union-attr] + + def test_ingestion_service_rebuild_skips_unchanged_sources(tmp_path: Path) -> None: allowed_root = tmp_path / "allowed" allowed_root.mkdir() From ebfaf5a468fbf6e8edca5152bdcdd392db9e5e1b Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 21:26:33 +0200 Subject: [PATCH 13/20] fix: update evals/run_evals.py for ragas 0.4.3 metric-instance API ragas 0.4.3 removed the pre-instantiated `ragas.metrics.collections` singletons from the classic `evaluate()` path; it now requires actual Metric instances (isinstance(m, Metric)), raising: TypeError: All metrics must be initialised metric objects, e.g: metrics=[BleuScore(), AspectCritic()] Import Faithfulness/AnswerRelevancy/ContextPrecision/ContextRecall from ragas.metrics (the legacy-compatible classes, still available via a deprecation shim) and construct instances for evaluate(). Also fixes a related break: EvaluationResult.__getitem__ now returns a list of per-row scores instead of a scalar, so float(result["metric"]) raised TypeError. Added _mean_score() to average per-row scores (NaN-safe), preserving the existing PASS_THRESHOLDS gating logic. --- evals/run_evals.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/evals/run_evals.py b/evals/run_evals.py index 53e85ae..9fca18c 100644 --- a/evals/run_evals.py +++ b/evals/run_evals.py @@ -18,6 +18,8 @@ import argparse import json +import math +import statistics import sys from datetime import UTC, datetime from pathlib import Path @@ -25,11 +27,11 @@ import httpx from datasets import Dataset from ragas import evaluate -from ragas.metrics.collections import ( - answer_relevancy, - context_precision, - context_recall, - faithfulness, +from ragas.metrics import ( + AnswerRelevancy, + ContextPrecision, + ContextRecall, + Faithfulness, ) DATASET_PATH = Path(__file__).parent / "dataset.json" @@ -93,6 +95,12 @@ def _build_hf_dataset( return Dataset.from_list(rows) +def _mean_score(values: list[float]) -> float: + """Average a metric's per-row scores, ignoring NaNs (ragas returns a list per key).""" + clean = [v for v in values if not math.isnan(v)] + return statistics.fmean(clean) if clean else math.nan + + def _print_summary(scores: dict[str, float]) -> bool: all_pass = True print("\n╔══════════════════════════════════╗") @@ -124,14 +132,14 @@ def main() -> None: print("Running RAGAS evaluation...") result = evaluate( dataset=dataset, - metrics=[faithfulness, answer_relevancy, context_precision, context_recall], + metrics=[Faithfulness(), AnswerRelevancy(), ContextPrecision(), ContextRecall()], ) scores: dict[str, float] = { - "faithfulness": float(result["faithfulness"]), - "answer_relevancy": float(result["answer_relevancy"]), - "context_precision": float(result["context_precision"]), - "context_recall": float(result["context_recall"]), + "faithfulness": _mean_score(result["faithfulness"]), + "answer_relevancy": _mean_score(result["answer_relevancy"]), + "context_precision": _mean_score(result["context_precision"]), + "context_recall": _mean_score(result["context_recall"]), } RESULTS_DIR.mkdir(parents=True, exist_ok=True) From 67af4f05989006aa5b986a8ffe3aaa84f4c6a6ec Mon Sep 17 00:00:00 2001 From: n0nuser Date: Tue, 7 Jul 2026 21:26:58 +0200 Subject: [PATCH 14/20] ci: gate RAGAS evals on pull requests touching rag/ingestion --- .github/workflows/evals-pr.yml | 39 ++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 .github/workflows/evals-pr.yml diff --git a/.github/workflows/evals-pr.yml b/.github/workflows/evals-pr.yml new file mode 100644 index 0000000..71f7fd2 --- /dev/null +++ b/.github/workflows/evals-pr.yml @@ -0,0 +1,39 @@ +name: RAGAS Evals (PR gate) + +on: + pull_request: + paths: + - "localrag/rag/**" + - "localrag/ingestion/**" + - "evals/**" + - "pyproject.toml" + +permissions: + contents: read + +jobs: + evals: + name: RAGAS evaluation (offline, blocking) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: astral-sh/setup-uv@v6 + with: + version: "latest" + + - name: Install dependencies + run: uv sync --all-extras + + - name: Run RAGAS evals (offline — no live API required) + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + run: uv run python evals/run_evals.py --offline + + - name: Upload eval results + if: always() + uses: actions/upload-artifact@v4 + with: + name: eval-results-pr-${{ github.event.pull_request.number }} + path: evals/results/ + if-no-files-found: warn From bb1780220cee2693c0e45a7adad3a4b97976e5b4 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 13:52:30 +0200 Subject: [PATCH 15/20] fix: ragas 0.4.3 compat + offline Ollama judge for evals/run_evals.py Modern ragas.metrics.collections classes require per-sample async .ascore() calls, not the deprecated singleton-style ragas.metrics imports; the judge LLM/embeddings now run against the same local Ollama instance LocalRAG itself uses (via Ollama's OpenAI-compatible /v1 endpoint and the openai client already a core dependency), so evals stay fully offline with no OPENAI_API_KEY required anywhere, including CI. --- .github/workflows/evals-pr.yml | 14 ++++- evals/run_evals.py | 100 +++++++++++++++++++++++++-------- 2 files changed, 88 insertions(+), 26 deletions(-) diff --git a/.github/workflows/evals-pr.yml b/.github/workflows/evals-pr.yml index 71f7fd2..55e6ff5 100644 --- a/.github/workflows/evals-pr.yml +++ b/.github/workflows/evals-pr.yml @@ -15,6 +15,11 @@ jobs: evals: name: RAGAS evaluation (offline, blocking) runs-on: ubuntu-latest + services: + ollama: + image: ollama/ollama:latest + ports: + - 11434:11434 steps: - uses: actions/checkout@v4 @@ -25,9 +30,12 @@ jobs: - name: Install dependencies run: uv sync --all-extras - - name: Run RAGAS evals (offline — no live API required) - env: - OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + - name: Pull judge models + run: | + curl -s http://localhost:11434/api/pull -d '{"name": "gemma3:4b"}' + curl -s http://localhost:11434/api/pull -d '{"name": "nomic-embed-text"}' + + - name: Run RAGAS evals (offline — judge runs against local Ollama, no external API key needed) run: uv run python evals/run_evals.py --offline - name: Upload eval results diff --git a/evals/run_evals.py b/evals/run_evals.py index 9fca18c..1a7027a 100644 --- a/evals/run_evals.py +++ b/evals/run_evals.py @@ -8,15 +8,24 @@ uv run python evals/run_evals.py [--api-url URL] [--api-key KEY] [--offline] Options: - --api-url LocalRAG API base URL (default: http://localhost:8000) - --api-key X-API-Key header value (empty = no auth) - --offline Skip live API calls; use stored contexts from dataset only - (requires ground-truth answers to already be in the dataset) + --api-url LocalRAG API base URL (default: http://localhost:8000) + --api-key X-API-Key header value (empty = no auth) + --offline Skip live API calls; use stored contexts from dataset only + (requires ground-truth answers to already be in the dataset) + --judge-model Ollama model used as the RAGAS LLM judge (default: gemma3:4b) + --ollama-url Ollama base URL for the judge/embeddings (default: http://localhost:11434) + +The RAGAS judge LLM and embeddings run on the same local Ollama instance +LocalRAG itself uses, via Ollama's OpenAI-compatible `/v1` endpoint and the +`openai` client already a core dependency of this project — no LangChain, no +new dependency, no external API key required. This matches LocalRAG's +offline-first positioning. """ from __future__ import annotations import argparse +import asyncio import json import math import statistics @@ -25,9 +34,10 @@ from pathlib import Path import httpx -from datasets import Dataset -from ragas import evaluate -from ragas.metrics import ( +from openai import AsyncOpenAI +from ragas.embeddings import OpenAIEmbeddings +from ragas.llms import llm_factory +from ragas.metrics.collections import ( AnswerRelevancy, ContextPrecision, ContextRecall, @@ -64,12 +74,12 @@ def _query_api(question: str, api_url: str, api_key: str) -> tuple[str, list[str return answer, contexts -def _build_hf_dataset( +def _build_rows( records: list[dict], api_url: str, api_key: str, offline: bool, -) -> Dataset: +) -> list[dict]: rows: list[dict] = [] for rec in records: question = rec["question"] @@ -92,15 +102,54 @@ def _build_hf_dataset( "ground_truth": ground_truth, } ) - return Dataset.from_list(rows) + return rows def _mean_score(values: list[float]) -> float: - """Average a metric's per-row scores, ignoring NaNs (ragas returns a list per key).""" + """Average a metric's per-row scores, ignoring NaNs.""" clean = [v for v in values if not math.isnan(v)] return statistics.fmean(clean) if clean else math.nan +async def _score_rows( + rows: list[dict], + *, + faithfulness: Faithfulness, + answer_relevancy: AnswerRelevancy, + context_precision: ContextPrecision, + context_recall: ContextRecall, +) -> dict[str, list[float]]: + per_metric: dict[str, list[float]] = { + "faithfulness": [], + "answer_relevancy": [], + "context_precision": [], + "context_recall": [], + } + for i, row in enumerate(rows, start=1): + print(f" scoring {i}/{len(rows)}: {row['question'][:60]}...") + user_input = row["question"] + response = row["answer"] + retrieved_contexts = row["contexts"] + reference = row["ground_truth"] + + faithfulness_result = await faithfulness.ascore( + user_input=user_input, response=response, retrieved_contexts=retrieved_contexts + ) + relevancy_result = await answer_relevancy.ascore(user_input=user_input, response=response) + precision_result = await context_precision.ascore( + user_input=user_input, reference=reference, retrieved_contexts=retrieved_contexts + ) + recall_result = await context_recall.ascore( + user_input=user_input, retrieved_contexts=retrieved_contexts, reference=reference + ) + + per_metric["faithfulness"].append(float(faithfulness_result.value)) + per_metric["answer_relevancy"].append(float(relevancy_result.value)) + per_metric["context_precision"].append(float(precision_result.value)) + per_metric["context_recall"].append(float(recall_result.value)) + return per_metric + + def _print_summary(scores: dict[str, float]) -> bool: all_pass = True print("\n╔══════════════════════════════════╗") @@ -121,26 +170,31 @@ def main() -> None: parser.add_argument("--api-url", default="http://localhost:8000") parser.add_argument("--api-key", default="") parser.add_argument("--offline", action="store_true") + parser.add_argument("--judge-model", default="gemma3:4b") + parser.add_argument("--ollama-url", default="http://localhost:11434") args = parser.parse_args() records: list[dict] = json.loads(DATASET_PATH.read_text(encoding="utf-8")) print(f"Loaded {len(records)} evaluation examples from {DATASET_PATH}") print("Building dataset" + (" (offline mode)" if args.offline else " (live API)") + "...") - dataset = _build_hf_dataset(records, args.api_url, args.api_key, offline=args.offline) - - print("Running RAGAS evaluation...") - result = evaluate( - dataset=dataset, - metrics=[Faithfulness(), AnswerRelevancy(), ContextPrecision(), ContextRecall()], + rows = _build_rows(records, args.api_url, args.api_key, offline=args.offline) + + print(f"Running RAGAS evaluation (judge={args.judge_model} via {args.ollama_url})...") + ollama_client = AsyncOpenAI(base_url=f"{args.ollama_url.rstrip('/')}/v1", api_key="ollama") + judge_llm = llm_factory(args.judge_model, client=ollama_client, adapter="instructor") + judge_embeddings = OpenAIEmbeddings(client=ollama_client, model="nomic-embed-text") + per_metric = asyncio.run( + _score_rows( + rows, + faithfulness=Faithfulness(llm=judge_llm), + answer_relevancy=AnswerRelevancy(llm=judge_llm, embeddings=judge_embeddings), + context_precision=ContextPrecision(llm=judge_llm), + context_recall=ContextRecall(llm=judge_llm), + ) ) - scores: dict[str, float] = { - "faithfulness": _mean_score(result["faithfulness"]), - "answer_relevancy": _mean_score(result["answer_relevancy"]), - "context_precision": _mean_score(result["context_precision"]), - "context_recall": _mean_score(result["context_recall"]), - } + scores: dict[str, float] = {name: _mean_score(values) for name, values in per_metric.items()} RESULTS_DIR.mkdir(parents=True, exist_ok=True) ts = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") From f8b6795fd962f2c66d6b7bfdeed5b29f321b6413 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 13:54:03 +0200 Subject: [PATCH 16/20] feat: background job registry for async directory ingest Claude-Session: https://claude.ai/code/session_01Md71iujmpvewiisGX5u33H --- docs/agent-navigation.md | 1 + docs/architecture.md | 3 +- localrag/api/dependencies.py | 6 +++ localrag/api/jobs.py | 85 ++++++++++++++++++++++++++++++++++ localrag/api/routers/ingest.py | 37 ++++++++++++++- localrag/api/schemas.py | 14 ++++++ localrag/api/service.py | 42 +++++++++++++++++ tests/test_api_ingest_async.py | 70 ++++++++++++++++++++++++++++ tests/test_api_jobs.py | 46 ++++++++++++++++++ 9 files changed, 302 insertions(+), 2 deletions(-) create mode 100644 localrag/api/jobs.py create mode 100644 tests/test_api_ingest_async.py create mode 100644 tests/test_api_jobs.py diff --git a/docs/agent-navigation.md b/docs/agent-navigation.md index 76eac4e..6beb482 100644 --- a/docs/agent-navigation.md +++ b/docs/agent-navigation.md @@ -31,6 +31,7 @@ Agents (and humans) move faster when they: | API app factory (lifespan, middleware, error handlers) | `localrag/api/main.py` | | HTTP ingest path validation (`INGEST_ROOTS`, URL decode) | `localrag/api/service.py`, `localrag/settings.py` (`is_path_allowed`), `localrag/api/exceptions.py` + `main.py` handler | | HTTP multipart file upload ingest (`POST /ingest/upload`) | `localrag/api/routers/ingest.py` (`ingest_upload`, Swagger limitations in `_UPLOAD_DESCRIPTION`), `localrag/api/service.py` (`ingest_upload`, `_stream_upload_to_disk`), `UPLOAD_DIR` / `UPLOAD_MAX_BYTES` in `localrag/settings.py` | +| Background ingest jobs | `localrag/api/jobs.py`, `localrag/api/routers/ingest.py` (async routes), `localrag/api/service.py` (`ingest_directory_async`, `get_ingest_job`) | | DI / shared service instances | `localrag/api/dependencies.py` | | Log format, levels, request ID | `localrag/logging_config.py`, `localrag/api/middleware.py`, `LOG_LEVEL` in `localrag/settings.py` | | API key auth | `localrag/api/dependencies.py` (`require_api_key`), `API_KEY` in `localrag/settings.py` | diff --git a/docs/architecture.md b/docs/architecture.md index 2730709..45d1788 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -42,7 +42,7 @@ flowchart LR F --> P --> LLM ``` -- **Ingest:** files → `loader` / `ingestion/parsers/*` → text → structural chunker (`localrag/ingestion/structural_chunker.py`) or fixed fallback (`localrag/ingestion/chunker.py`) → `OllamaEmbedder` → `VectorStore` (Chroma, persistent path from settings). The **HTTP** ingest flow runs path decode, existence checks, and `INGEST_ROOTS` in `localrag/api/service.py` (`ingest_file` / `ingest_directory`), then calls `IngestionService`; optional per-request `embed_model` overrides `OLLAMA_EMBED_MODEL`. Failures raise `IngestApiError` → JSON in `main.py`. CLI ingests call `IngestionService` directly. `POST /ingest/upload` (`ingest_upload` in `service.py`) takes a multipart file instead of a server path: it validates the extension against `loader.SUPPORTED_EXTENSIONS`, streams it to disk under `UPLOAD_DIR` in 1 MiB chunks while enforcing `UPLOAD_MAX_BYTES` (bypassing `INGEST_ROOTS`, since the server picks the destination), then calls `IngestionService.ingest_file` the same way. See the endpoint's OpenAPI description for upload limitations (no AV scan, extension-only validation, single file per request). +- **Ingest:** files → `loader` / `ingestion/parsers/*` → text → structural chunker (`localrag/ingestion/structural_chunker.py`) or fixed fallback (`localrag/ingestion/chunker.py`) → `OllamaEmbedder` → `VectorStore` (Chroma, persistent path from settings). The **HTTP** ingest flow runs path decode, existence checks, and `INGEST_ROOTS` in `localrag/api/service.py` (`ingest_file` / `ingest_directory`), then calls `IngestionService`; optional per-request `embed_model` overrides `OLLAMA_EMBED_MODEL`. Failures raise `IngestApiError` → JSON in `main.py`. CLI ingests call `IngestionService` directly. `POST /ingest/upload` (`ingest_upload` in `service.py`) takes a multipart file instead of a server path: it validates the extension against `loader.SUPPORTED_EXTENSIONS`, streams it to disk under `UPLOAD_DIR` in 1 MiB chunks while enforcing `UPLOAD_MAX_BYTES` (bypassing `INGEST_ROOTS`, since the server picks the destination), then calls `IngestionService.ingest_file` the same way. See the endpoint's OpenAPI description for upload limitations (no AV scan, extension-only validation, single file per request). For long-running directory ingests, `POST /ingest/directory/async` (`ingest_directory_async` in `service.py`) runs the same path validation synchronously, then submits the actual `IngestionService.ingest_directory` call to the in-process `JobRegistry` (`localrag/api/jobs.py`) and returns `202 {job_id, status: "pending"}` immediately; poll `GET /ingest/jobs/{job_id}` (`get_ingest_job`) for `running` / `done` (with `result`) / `failed` (with `error`). - **Rebuild:** `POST /collections/rebuild` and `localrag collections rebuild` list distinct `source` values in the active collection, drop vectors for missing files, and re-chunk/re-embed remaining paths (optional `embed_model` override). Implemented in `IngestionService.rebuild_collection`. - **Query (JSON):** `POST /query` returns a complete `QueryResponse` (answer, sources, latency_ms, model) from `query_json` in `localrag/api/service.py`. Retrieval supports vector-only and hybrid (vector + BM25 with reciprocal-rank fusion), then applies optional freshness decay based on chunk `ingested_at`. `RAGEngine` generates the answer via its injected `provider` (a `BaseLLMProvider` built by `llm/factory.py::build_provider`, resilience-wrapped), so `LLM_BACKEND` genuinely governs which backend answers `/query` — it is no longer hard-wired to Ollama. Requires `X-API-Key` when `API_KEY` is set. - **Query (SSE stream):** `POST /query/stream` streams tokens as Server-Sent Events. Retrieval runs synchronously first (`get_query_contexts`) so errors map to HTTP before SSE starts, then tokens are mapped via `iter_query_sse_events`. Token streaming likewise goes through `RAGEngine.provider.stream_from_prompt(...)`, so `LLM_BACKEND` governs the streaming path too. @@ -61,6 +61,7 @@ flowchart LR | API key auth | `localrag/api/dependencies.py` | `require_api_key` dependency — enforces `X-API-Key` when `API_KEY` env var is set | | Prometheus metrics | `localrag/api/routers/metrics.py` | `GET /metrics` via `prometheus_client.generate_latest()` | | HTTP API (persistence) | `localrag/api/repository.py` | `ChromaCollectionRepository` → `VectorStore` for collection list/delete and health’s collection list | +| Background ingest jobs | `localrag/api/jobs.py` | `JobRegistry` — in-memory `ThreadPoolExecutor`-backed job store for `POST /ingest/directory/async`; no persistence across process restarts | | CLI | `localrag/cli/app.py`, `localrag/cli/commands/*` | `localrag` Typer entry (`pyproject` `[project.scripts]`) | | Ingestion orchestration | `localrag/ingestion/service.py` | `IngestionService`: paths → parse → chunk → embed → upsert | | File formats | `localrag/ingestion/parsers/*` | pdf (with OCR fallback via pypdfium2 + pytesseract), docx, markdown, text, code | diff --git a/localrag/api/dependencies.py b/localrag/api/dependencies.py index 950b1fc..2ce136e 100644 --- a/localrag/api/dependencies.py +++ b/localrag/api/dependencies.py @@ -6,6 +6,7 @@ from fastapi import Depends, HTTPException, Security from fastapi.security import APIKeyHeader +from localrag.api.jobs import JobRegistry from localrag.api.repository import ChromaCollectionRepository from localrag.ingestion.embedder import OllamaEmbedder from localrag.ingestion.service import IngestionService @@ -82,6 +83,11 @@ def get_bm25_index() -> Bm25Index: return Bm25Index.from_vector_store(get_vector_store()) +@lru_cache(maxsize=1) +def get_job_registry() -> JobRegistry: + return JobRegistry() + + def require_api_key( key: str | None = Security(_api_key_header), settings: Settings = Depends(get_settings), diff --git a/localrag/api/jobs.py b/localrag/api/jobs.py new file mode 100644 index 0000000..d752639 --- /dev/null +++ b/localrag/api/jobs.py @@ -0,0 +1,85 @@ +"""In-memory background job registry for long-running ingest operations. + +Jobs live only for the lifetime of the API process — there is no persistence +across restarts. This matches LocalRAG's single-process, offline-first scope; +a durable job queue backed by an external broker is intentionally out of scope. +""" + +from __future__ import annotations + +import logging +import threading +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from datetime import UTC, datetime +from enum import StrEnum +from typing import Any +from uuid import uuid4 + +logger = logging.getLogger(__name__) + + +class JobStatus(StrEnum): + PENDING = "pending" + RUNNING = "running" + DONE = "done" + FAILED = "failed" + + +@dataclass +class Job: + job_id: str + status: JobStatus + created_at: str + result: dict[str, Any] | None = None + error: str | None = None + + +@dataclass +class JobRegistry: + _jobs: dict[str, Job] = field(default_factory=dict) + _lock: threading.Lock = field(default_factory=threading.Lock) + _executor: ThreadPoolExecutor = field( + default_factory=lambda: ThreadPoolExecutor(max_workers=2, thread_name_prefix="ingest-job") + ) + + def submit(self, work: Callable[[], dict[str, Any]]) -> str: + job_id = uuid4().hex + job = Job( + job_id=job_id, + status=JobStatus.PENDING, + created_at=datetime.now(UTC).isoformat(), + ) + with self._lock: + self._jobs[job_id] = job + self._executor.submit(self._run, job_id, work) + return job_id + + def _run(self, job_id: str, work: Callable[[], dict[str, Any]]) -> None: + with self._lock: + self._jobs[job_id].status = JobStatus.RUNNING + try: + result = work() + except Exception as exc: # captured for the caller, not re-raised in a background thread + logger.exception("ingest_job_failed job_id=%s", job_id) + with self._lock: + self._jobs[job_id].status = JobStatus.FAILED + self._jobs[job_id].error = str(exc) + return + with self._lock: + self._jobs[job_id].status = JobStatus.DONE + self._jobs[job_id].result = result + + def get(self, job_id: str) -> Job | None: + with self._lock: + job = self._jobs.get(job_id) + if job is None: + return None + return Job( + job_id=job.job_id, + status=job.status, + created_at=job.created_at, + result=job.result, + error=job.error, + ) diff --git a/localrag/api/routers/ingest.py b/localrag/api/routers/ingest.py index c18e9e5..4375f0c 100644 --- a/localrag/api/routers/ingest.py +++ b/localrag/api/routers/ingest.py @@ -3,12 +3,20 @@ from fastapi import APIRouter, Depends, File, Form, UploadFile from localrag.api import service as api_service -from localrag.api.dependencies import get_api_settings, get_ingestion_service, require_api_key +from localrag.api.dependencies import ( + get_api_settings, + get_ingestion_service, + get_job_registry, + require_api_key, +) +from localrag.api.jobs import JobRegistry from localrag.api.schemas import ( IngestDirectoryRequest, IngestDirectoryResponse, IngestFileRequest, IngestFileResponse, + IngestJobResponse, + IngestJobStatusResponse, ) from localrag.ingestion.service import IngestionService from localrag.settings import Settings @@ -86,3 +94,30 @@ def ingest_directory( ingestion_service: IngestionService = Depends(get_ingestion_service), ) -> IngestDirectoryResponse: return api_service.ingest_directory(request, settings, ingestion_service) + + +@router.post( + "/ingest/directory/async", + response_model=IngestJobResponse, + status_code=202, + summary="Ingest a directory in the background", + description=( + "Submits a directory ingest to an in-process background job (no external broker; " + "job state is lost on process restart). Poll GET /ingest/jobs/{job_id} for status." + ), +) +def ingest_directory_async( + request: IngestDirectoryRequest, + settings: Settings = Depends(get_api_settings), + ingestion_service: IngestionService = Depends(get_ingestion_service), + job_registry: JobRegistry = Depends(get_job_registry), +) -> IngestJobResponse: + return api_service.ingest_directory_async(request, settings, ingestion_service, job_registry) + + +@router.get("/ingest/jobs/{job_id}", response_model=IngestJobStatusResponse) +def get_ingest_job( + job_id: str, + job_registry: JobRegistry = Depends(get_job_registry), +) -> IngestJobStatusResponse: + return api_service.get_ingest_job(job_id, job_registry) diff --git a/localrag/api/schemas.py b/localrag/api/schemas.py index cf146d1..a1521d7 100644 --- a/localrag/api/schemas.py +++ b/localrag/api/schemas.py @@ -189,6 +189,20 @@ class IngestDirectoryResponse(BaseModel): ) +class IngestJobResponse(BaseModel): + job_id: str = Field(description="Opaque identifier to poll via GET /ingest/jobs/{job_id}.") + status: str = Field(description="Initial job status.", examples=["pending"]) + + +class IngestJobStatusResponse(BaseModel): + job_id: str + status: str = Field(examples=["pending", "running", "done", "failed"]) + result: dict[str, object] | None = Field( + default=None, description="Ingest result payload once status is 'done'." + ) + error: str | None = Field(default=None, description="Error message when status is 'failed'.") + + class HealthResponse(BaseModel): status: str = Field( description="Overall API process status.", diff --git a/localrag/api/service.py b/localrag/api/service.py index 537301f..d659833 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -14,6 +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.repository import ChromaCollectionRepository from localrag.api.schemas import ( CollectionDeleteResponse, @@ -24,6 +25,8 @@ IngestDirectoryResponse, IngestFileRequest, IngestFileResponse, + IngestJobResponse, + IngestJobStatusResponse, QueryRequest, QueryResponse, RebuildCollectionRequest, @@ -205,6 +208,45 @@ def ingest_directory( ) +def ingest_directory_async( + request: IngestDirectoryRequest, + settings: Settings, + ingestion_service: IngestionService, + job_registry: JobRegistry, +) -> IngestJobResponse: + path = path_from_ingest_request(request.path).resolve() + if not path.is_dir(): + raise IngestApiError(HTTPStatus.BAD_REQUEST, "Path must be an existing directory.") + if not is_path_allowed(path, settings.ingest_roots): + raise IngestApiError(HTTPStatus.FORBIDDEN, "Path is not under configured ingest roots.") + + def work() -> dict[str, Any]: + result = ingestion_service.ingest_directory( + path, recursive=request.recursive, embed_model=request.embed_model + ) + return { + "status": "ok", + "files_processed": result.files_processed, + "total_chunks": result.total_chunks, + "failed_sources": [ + {"source": f.source, "error": f.error} for f in result.failed_sources + ], + } + + job_id = job_registry.submit(work) + logger.info("ingest_directory_async_submitted job_id=%s path=%s", job_id, path) + return IngestJobResponse(job_id=job_id, status="pending") + + +def get_ingest_job(job_id: str, job_registry: JobRegistry) -> IngestJobStatusResponse: + job = job_registry.get(job_id) + if job is None: + raise IngestApiError(HTTPStatus.NOT_FOUND, f"Unknown job id '{job_id}'.") + return IngestJobStatusResponse( + job_id=job.job_id, status=job.status.value, result=job.result, error=job.error + ) + + def ingest_upload( file_name: str, file_obj: BinaryIO, diff --git a/tests/test_api_ingest_async.py b/tests/test_api_ingest_async.py new file mode 100644 index 0000000..efc256e --- /dev/null +++ b/tests/test_api_ingest_async.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import time +from dataclasses import dataclass +from pathlib import Path + +from fastapi.testclient import TestClient + +from localrag.api.dependencies import get_api_settings, get_ingestion_service, get_job_registry +from localrag.api.jobs import JobRegistry +from localrag.api.main import app +from localrag.ingestion.service import IngestionResult +from localrag.settings import Settings + + +@dataclass +class StubIngestionService: + result: IngestionResult + + def ingest_directory( + self, path: Path, recursive: bool | None = None, embed_model: str | None = None + ) -> IngestionResult: + return self.result + + +def test_ingest_directory_async_returns_job_id_then_reports_done(tmp_path: Path) -> None: + allowed_root = tmp_path / "allowed" + allowed_root.mkdir() + + settings = Settings(ingest_roots=[str(allowed_root)]) + ingestion = StubIngestionService( + result=IngestionResult(files_processed=1, total_chunks=4, processed_sources=[]) + ) + registry = JobRegistry() + + app.dependency_overrides[get_api_settings] = lambda: settings + app.dependency_overrides[get_ingestion_service] = lambda: ingestion + app.dependency_overrides[get_job_registry] = lambda: registry + client = TestClient(app) + + submitted = client.post("/ingest/directory/async", json={"path": str(allowed_root)}) + assert submitted.status_code == 202 + job_id = submitted.json()["job_id"] + assert submitted.json()["status"] == "pending" + + deadline = time.monotonic() + 5.0 + status_body: dict[str, object] = {} + while time.monotonic() < deadline: + polled = client.get(f"/ingest/jobs/{job_id}") + assert polled.status_code == 200 + status_body = polled.json() + if status_body["status"] == "done": + break + time.sleep(0.01) + + assert status_body["status"] == "done" + assert status_body["result"]["total_chunks"] == 4 # type: ignore[index] + + app.dependency_overrides.clear() + + +def test_ingest_job_status_unknown_id_returns_404() -> None: + registry = JobRegistry() + app.dependency_overrides[get_job_registry] = lambda: registry + client = TestClient(app) + + response = client.get("/ingest/jobs/does-not-exist") + assert response.status_code == 404 + + app.dependency_overrides.clear() diff --git a/tests/test_api_jobs.py b/tests/test_api_jobs.py new file mode 100644 index 0000000..3a7c000 --- /dev/null +++ b/tests/test_api_jobs.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import time + +from localrag.api.jobs import Job, JobRegistry, JobStatus + + +def _wait_for_terminal(registry: JobRegistry, job_id: str, timeout: float = 5.0) -> Job | None: + deadline = time.monotonic() + timeout + job = registry.get(job_id) + while job is not None and job.status in (JobStatus.PENDING, JobStatus.RUNNING): + if time.monotonic() > deadline: + raise AssertionError("job did not finish in time") + time.sleep(0.01) + job = registry.get(job_id) + return job + + +def test_job_registry_submit_runs_work_in_background_and_reports_done() -> None: + registry = JobRegistry() + job_id = registry.submit(lambda: {"value": 42}) + + job = _wait_for_terminal(registry, job_id) + + assert job is not None + assert job.status == JobStatus.DONE + assert job.result == {"value": 42} + + +def test_job_registry_marks_failed_on_exception() -> None: + registry = JobRegistry() + + def boom() -> dict[str, object]: + raise RuntimeError("nope") + + job_id = registry.submit(boom) + job = _wait_for_terminal(registry, job_id) + + assert job is not None + assert job.status == JobStatus.FAILED + assert job.error == "nope" + + +def test_job_registry_get_returns_none_for_unknown_job() -> None: + registry = JobRegistry() + assert registry.get("nope") is None From ec9e2816ac910d397fe2088d4c70dc82251cd357 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 13:54:27 +0200 Subject: [PATCH 17/20] feat: in-process TTL cache for repeated queries Claude-Session: https://claude.ai/code/session_01Md71iujmpvewiisGX5u33H --- .env.example | 4 ++ docs/rag-retrieval.md | 33 +++++++++++++ localrag/api/dependencies.py | 9 ++++ localrag/api/routers/query.py | 6 ++- localrag/api/service.py | 17 ++++++- localrag/rag/query_cache.py | 42 +++++++++++++++++ localrag/settings.py | 4 ++ pyproject.toml | 1 + tests/test_api_main_exceptions.py | 3 +- tests/test_api_query_cache.py | 78 +++++++++++++++++++++++++++++++ tests/test_query_cache.py | 34 ++++++++++++++ uv.lock | 11 +++++ 12 files changed, 238 insertions(+), 4 deletions(-) create mode 100644 localrag/rag/query_cache.py create mode 100644 tests/test_api_query_cache.py create mode 100644 tests/test_query_cache.py diff --git a/.env.example b/.env.example index e7bafc3..bb5485c 100644 --- a/.env.example +++ b/.env.example @@ -54,6 +54,10 @@ RERANK_ENABLED=false RERANK_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2 RERANK_FETCH_K=20 +# In-process TTL cache for repeated/near-identical queries (0 disables; no external cache). +QUERY_CACHE_TTL_SECONDS=0 +QUERY_CACHE_MAXSIZE=256 + API_HOST=0.0.0.0 API_PORT=8000 diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index a5138e2..0be28f2 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -187,6 +187,39 @@ a security boundary — anyone with API access can still query across all per-tenant isolation is ever required, revisit as a dedicated (out-of-scope-for-this-plan) access-control project. +## Query caching (optional) + +Disabled by default (`QUERY_CACHE_TTL_SECONDS=0`). When enabled (set a positive +TTL in seconds), `POST /query` (`localrag/api/routers/query.py`) is served +through an in-process `QueryCache` (`localrag/rag/query_cache.py`), wired via +`localrag/api/dependencies.py::get_query_cache` (an `lru_cache`-memoized +singleton, so all requests within one process share the same cache) and +passed into `localrag.api.service.query_json`. + +Cache keys are an exact-match SHA-256 hash over the normalized question +(stripped, lowercased), `model`, `n_results`, and `retrieval_mode` +(`make_cache_key`) — this is not semantic/fuzzy matching, so any change to +wording, model, result count, or retrieval mode is a cache miss. Cached values +are the full serialized `QueryResponse` (`response.model_dump()`), so a cache +hit replays `answer`, `sources` (including `heading_path`/`chunk_type`), +`latency_ms`, `model`, and `low_confidence` exactly as they were on the +original request — a cached low-confidence refusal is served back as +low-confidence, not silently upgraded. + +`QUERY_CACHE_MAXSIZE` bounds the number of entries (least-recently-used +eviction via `cachetools.TTLCache`) independent of the TTL. + +This cache is in-process only — it is **not** shared across multiple +`uvicorn` worker processes (each worker gets its own `QueryCache` instance), +so cache hit rate degrades as worker count increases. If LocalRAG ever runs +multi-process/multi-replica in production, a shared external cache (e.g. +Redis) would be the upgrade path; that is out of scope for the current +single-process deployment model. + +Cache hits are not separately audit-logged — a served-from-cache response +still only produces the `query_cache_hit` log line, since it bypasses the +retriever and LLM call entirely. + ## Ingestion metadata dependencies Freshness and debugging depend on chunk metadata written during ingestion: diff --git a/localrag/api/dependencies.py b/localrag/api/dependencies.py index 2ce136e..2da8c34 100644 --- a/localrag/api/dependencies.py +++ b/localrag/api/dependencies.py @@ -13,6 +13,7 @@ from localrag.llm.factory import build_provider from localrag.rag.bm25_index import Bm25Index from localrag.rag.engine import RAGEngine +from localrag.rag.query_cache import QueryCache from localrag.rag.reranker import CrossEncoderReranker from localrag.rag.retriever import Retriever from localrag.settings import Settings, get_settings @@ -67,6 +68,14 @@ def get_engine() -> RAGEngine: ) +@lru_cache(maxsize=1) +def get_query_cache() -> QueryCache: + settings = get_settings() + return QueryCache( + maxsize=settings.query_cache_maxsize, ttl_seconds=settings.query_cache_ttl_seconds + ) + + @lru_cache(maxsize=1) def get_ingestion_service() -> IngestionService: settings = get_settings() diff --git a/localrag/api/routers/query.py b/localrag/api/routers/query.py index 63ec236..5a338ce 100644 --- a/localrag/api/routers/query.py +++ b/localrag/api/routers/query.py @@ -4,9 +4,10 @@ from sse_starlette import EventSourceResponse from localrag.api import service as api_service -from localrag.api.dependencies import get_engine, require_api_key +from localrag.api.dependencies import get_engine, get_query_cache, require_api_key from localrag.api.schemas import QueryRequest, QueryResponse from localrag.rag.engine import RAGEngine +from localrag.rag.query_cache import QueryCache router = APIRouter(prefix="", tags=["query"], dependencies=[Depends(require_api_key)]) @@ -15,9 +16,10 @@ def query( request: QueryRequest, engine: RAGEngine = Depends(get_engine), + query_cache: QueryCache = Depends(get_query_cache), ) -> QueryResponse: """Retrieve context and return a complete JSON response with answer, sources, and latency.""" - return api_service.query_json(request, engine) + return api_service.query_json(request, engine, query_cache) @router.post("/query/stream", summary="Query (SSE stream)") diff --git a/localrag/api/service.py b/localrag/api/service.py index d659833..515ca61 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -40,6 +40,7 @@ from localrag.ollama.schemas import OllamaTagsResponse, parse_ollama_json from localrag.rag.engine import RAGEngine from localrag.rag.exceptions import RetrievalError +from localrag.rag.query_cache import QueryCache, make_cache_key from localrag.settings import Settings, is_path_allowed logger = logging.getLogger(__name__) @@ -317,9 +318,21 @@ def _unique_upload_destination(directory: Path, file_name: str) -> Path: return directory / f"{stem}-{uuid4().hex[:8]}{suffix}" -def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: +def query_json( + request: QueryRequest, engine: RAGEngine, query_cache: QueryCache | None = None +) -> QueryResponse: """Blocking JSON query — retrieves context then generates a full answer.""" t0 = time.perf_counter() + cache_key: str | None = None + if query_cache is not None: + cache_key = make_cache_key( + request.question, request.model, request.n_results, engine.settings.retrieval_mode + ) + cached = query_cache.get(cache_key) + if cached is not None: + logger.info("query_cache_hit") + return QueryResponse(**cached) + try: contexts = engine.retriever.retrieve( question=request.question, @@ -380,6 +393,8 @@ def query_json(request: QueryRequest, engine: RAGEngine) -> QueryResponse: model=used_model, latency_ms=latency_ms, ) + if query_cache is not None and cache_key is not None: + query_cache.set(cache_key, response.model_dump()) return response diff --git a/localrag/rag/query_cache.py b/localrag/rag/query_cache.py new file mode 100644 index 0000000..89989f5 --- /dev/null +++ b/localrag/rag/query_cache.py @@ -0,0 +1,42 @@ +"""In-process TTL cache for repeated/near-identical queries (no external service).""" + +from __future__ import annotations + +import hashlib +import json +from typing import Any + +from cachetools import TTLCache + + +def make_cache_key( + question: str, model: str | None, n_results: int | None, retrieval_mode: str +) -> str: + payload = json.dumps( + { + "question": question.strip().lower(), + "model": model, + "n_results": n_results, + "retrieval_mode": retrieval_mode, + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest() + + +class QueryCache: + def __init__(self, maxsize: int, ttl_seconds: float) -> None: + self._enabled = ttl_seconds > 0 and maxsize > 0 + self._cache: TTLCache[str, dict[str, Any]] = TTLCache( + maxsize=max(1, maxsize), ttl=max(0.001, ttl_seconds) + ) + + def get(self, key: str) -> dict[str, Any] | None: + if not self._enabled: + return None + return self._cache.get(key) + + def set(self, key: str, value: dict[str, Any]) -> None: + if not self._enabled: + return + self._cache[key] = value diff --git a/localrag/settings.py b/localrag/settings.py index 1497335..e99c3a1 100644 --- a/localrag/settings.py +++ b/localrag/settings.py @@ -124,6 +124,10 @@ class Settings(BaseSettings): rerank_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2" rerank_fetch_k: int = 20 + # In-process TTL cache for repeated/near-identical queries (0 disables; no external cache). + query_cache_ttl_seconds: float = 0.0 + query_cache_maxsize: int = 256 + api_host: str = "0.0.0.0" # nosec B104 — configurable bind address, default intentional api_port: int = 8000 diff --git a/pyproject.toml b/pyproject.toml index 274ca97..6c6c0f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "python-multipart>=0.0.20", "tenacity>=9", "pybreaker>=1", + "cachetools>=5", ] [project.scripts] diff --git a/tests/test_api_main_exceptions.py b/tests/test_api_main_exceptions.py index c556689..d3b8686 100644 --- a/tests/test_api_main_exceptions.py +++ b/tests/test_api_main_exceptions.py @@ -14,7 +14,7 @@ ) from localrag.api.main import app from localrag.rag.exceptions import RetrievalError -from localrag.settings import Settings +from localrag.settings import Settings, get_settings @respx.mock @@ -72,6 +72,7 @@ def retrieve( @dataclass class FailingQueryEngine: retriever: FailingRetriever = field(default_factory=FailingRetriever) + settings: Settings = field(default_factory=get_settings) def test_query_maps_retrieval_failure_to_502() -> None: diff --git a/tests/test_api_query_cache.py b/tests/test_api_query_cache.py new file mode 100644 index 0000000..21437cf --- /dev/null +++ b/tests/test_api_query_cache.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from collections.abc import Generator +from dataclasses import dataclass, field +from typing import Any + +from fastapi.testclient import TestClient + +from localrag.api.dependencies import get_engine, get_query_cache +from localrag.api.main import app +from localrag.llm.providers.base import BaseLLMProvider +from localrag.llm.types import LLMResponse +from localrag.rag.engine import RAGEngine +from localrag.rag.query_cache import QueryCache +from localrag.settings import Settings + + +@dataclass +class CountingRetriever: + calls: list[str] + contexts: list[dict[str, object]] + + def retrieve( + self, + question: str, + n_results: int | None = None, + metadata_filter: dict[str, object] | None = None, + ) -> list[dict[str, object]]: + self.calls.append(question) + return self.contexts + + +@dataclass +class FakeProvider(BaseLLMProvider): + """Minimal stub standing in for the real LLM provider (see tests/test_rag_engine.py).""" + + tokens: list[str] = field(default_factory=lambda: ["answer"]) + + def generate(self, prompt: str, context: list[str], *, model: str | None = None) -> LLMResponse: + raise NotImplementedError + + def stream( + self, prompt: str, context: list[str], *, model: str | None = None + ) -> Generator[dict[str, Any]]: + raise NotImplementedError + + def generate_from_prompt(self, prompt: str, *, model: str | None = None) -> LLMResponse: + raise NotImplementedError + + def stream_from_prompt( + self, prompt: str, *, model: str | None = None + ) -> Generator[dict[str, Any]]: + for token in self.tokens: + yield {"type": "token", "token": token} + yield {"type": "final", "sources": []} + + def count_tokens(self, text: str) -> int: + return len(text.split()) + + +def test_repeat_query_is_served_from_cache_without_calling_retriever() -> None: + settings = Settings(ollama_base_url="http://ollama:11434", ollama_llm_model="llm") + retriever = CountingRetriever(calls=[], contexts=[]) + engine = RAGEngine(settings=settings, retriever=retriever, provider=FakeProvider()) # type: ignore[arg-type] + cache = QueryCache(maxsize=10, ttl_seconds=60) + + app.dependency_overrides[get_engine] = lambda: engine + app.dependency_overrides[get_query_cache] = lambda: cache + client = TestClient(app) + + first = client.post("/query", json={"question": "What is X?"}) + second = client.post("/query", json={"question": "What is X?"}) + + assert first.status_code == 200 + assert second.status_code == 200 + assert retriever.calls == ["What is X?"] # second request served from cache + + app.dependency_overrides.clear() diff --git a/tests/test_query_cache.py b/tests/test_query_cache.py new file mode 100644 index 0000000..c3cf846 --- /dev/null +++ b/tests/test_query_cache.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import time + +from localrag.rag.query_cache import QueryCache, make_cache_key + + +def test_query_cache_returns_none_when_disabled() -> None: + cache = QueryCache(maxsize=10, ttl_seconds=0) + key = make_cache_key("q", None, None, "hybrid") + cache.set(key, {"answer": "x"}) + assert cache.get(key) is None + + +def test_query_cache_hits_within_ttl_and_expires_after() -> None: + cache = QueryCache(maxsize=10, ttl_seconds=0.05) + key = make_cache_key("q", "m", 5, "hybrid") + cache.set(key, {"answer": "cached"}) + + assert cache.get(key) == {"answer": "cached"} + time.sleep(0.1) + assert cache.get(key) is None + + +def test_make_cache_key_is_case_and_whitespace_insensitive() -> None: + key_a = make_cache_key(" What is X? ", "m", 5, "hybrid") + key_b = make_cache_key("what is x?", "m", 5, "hybrid") + assert key_a == key_b + + +def test_make_cache_key_differs_by_model() -> None: + key_a = make_cache_key("q", "model-a", 5, "hybrid") + key_b = make_cache_key("q", "model-b", 5, "hybrid") + assert key_a != key_b diff --git a/uv.lock b/uv.lock index 0aee460..48bd058 100644 --- a/uv.lock +++ b/uv.lock @@ -304,6 +304,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0d/fe/6bea5c9162869c5beba5d9c8abbed835ec85bf1ec1fba05a3822325c45f3/build-1.5.0-py3-none-any.whl", hash = "sha256:13f3eecb844759ab66efec90ca17639bbf14dc06cb2fdf37a9010322d9c50a6f", size = 26018, upload-time = "2026-04-30T03:18:23.644Z" }, ] +[[package]] +name = "cachetools" +version = "7.1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f4/8b/0d3945a13955303b81272f759a0331e54c5c793da455e6f5706b89d2639c/cachetools-7.1.4.tar.gz", hash = "sha256:437f55a4e0c1b01a4f3077cc470e6991d47430970e36fbcb77e2be0df4fc1cd6", size = 40085, upload-time = "2026-05-21T22:40:43.376Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/7b/1fc1c09cc0756cf25861a3be10565915953876da48bb228fb9a672b20a42/cachetools-7.1.4-py3-none-any.whl", hash = "sha256:323dc4127934744db5b54eb4924482d7edafbf9554e820d1531c2e08c0e4ef54", size = 16761, upload-time = "2026-05-21T22:40:41.845Z" }, +] + [[package]] name = "certifi" version = "2026.5.20" @@ -1438,6 +1447,7 @@ version = "0.1.0" source = { editable = "." } dependencies = [ { name = "anthropic" }, + { name = "cachetools" }, { name = "charset-normalizer" }, { name = "chromadb" }, { name = "fastapi" }, @@ -1481,6 +1491,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "anthropic", specifier = ">=0.97.0" }, + { name = "cachetools", specifier = ">=5" }, { name = "charset-normalizer", specifier = ">=3.4" }, { name = "chromadb", specifier = ">=1.0" }, { name = "fastapi", specifier = ">=0.115" }, From cee607bbc33e71effa99d0456a13c7c84644ad0f Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 16:17:15 +0200 Subject: [PATCH 18/20] docs: rename AGENTS.md to CLAUDE.md, add session reasoning logs log_reasonings/ captures decisions made during the production-RAG-hardening plan rollout (ragas judge LLM choice, agent fan-out strategy, retry discipline) as lighter-weight session notes alongside the formal docs/adr/. --- AGENTS.md => CLAUDE.md | 20 ++++++ .../2026-07-07-agent-fanout-strategy.md | 44 ++++++++++++ .../2026-07-07-agent-retry-discipline.md | 46 +++++++++++++ .../2026-07-07-ragas-offline-judge.md | 67 +++++++++++++++++++ 4 files changed, 177 insertions(+) rename AGENTS.md => CLAUDE.md (64%) create mode 100644 log_reasonings/2026-07-07-agent-fanout-strategy.md create mode 100644 log_reasonings/2026-07-07-agent-retry-discipline.md create mode 100644 log_reasonings/2026-07-07-ragas-offline-judge.md diff --git a/AGENTS.md b/CLAUDE.md similarity index 64% rename from AGENTS.md rename to CLAUDE.md index 97fee43..7dcf054 100644 --- a/AGENTS.md +++ b/CLAUDE.md @@ -22,6 +22,26 @@ The FastAPI layer follows a **light domain-driven** split: Domain packages (`localrag/ingestion/`, `localrag/rag/`, `localrag/storage/`) keep their own services and types; the API service calls into them (e.g. `IngestionService`, `RAGEngine`). +## Claude Code skills for this repo + +If you're running as Claude Code, these installed skills/subagents map directly onto LocalRAG's stack—reach for them instead of re-deriving generic advice: + +| Area | Skill / agent | When it applies here | +| --- | --- | --- | +| Git workflow | `superpowers:using-git-worktrees`, `git-pr-workflows:git-workflow` | Isolating longer-lived `feat/…`/`fix/…` work, or orchestrating review → PR beyond the basic steps in `.cursor/skills/trunk-feature-workflow/SKILL.md` | +| API/DDD layering | `backend-development:architecture-patterns`, `backend-development:api-design-principles` | Reshaping the schemas/service/repository/router split above, or designing new routes | +| Python style & types | `python-development:python-code-style`, `python-development:python-type-safety` | Anything ruff/mypy already gate in `.pre-commit-config.yaml`—use before relying on defaults | +| Config & settings | `python-development:python-configuration` | Changes to `localrag/settings.py` / `.env.example` (pydantic-settings) | +| Error handling & resilience | `python-development:python-error-handling`, `python-development:python-resilience` | `HttpMappedError` subclasses, ingestion retry/batch logic | +| Observability | `python-development:python-observability` | structlog / Prometheus metrics work (`localrag/logging_config.py`, `api/routers/metrics.py`) | +| Testing | `python-development:python-testing-patterns`, `superpowers:test-driven-development` | Adding/extending `tests/` (pytest, pytest-asyncio, respx) | +| Debugging | `superpowers:systematic-debugging`, `diagnosing-bugs` | Any bug/regression—before proposing a fix | +| LLM/provider code | `claude-api` | Anthropic/OpenAI/Ollama provider work in `localrag/llm/` | +| Security | `security-review`, `backend-api-security:backend-security-coder` | Auth (`API_KEY`), path validation (`is_path_allowed`), upload handling, anything bandit flags | +| Verification | `verify` | Before claiming a change works—exercise the flow, don't just trust lint/tests | + +This is a pointer, not a guarantee of installation—confirm availability in your environment before relying on one. + ## Documentation maintenance for agents When you change anything that affects **how agents find or reason about the codebase**, update the relevant docs **in the same change** (same PR). At minimum: diff --git a/log_reasonings/2026-07-07-agent-fanout-strategy.md b/log_reasonings/2026-07-07-agent-fanout-strategy.md new file mode 100644 index 0000000..a2bbcf7 --- /dev/null +++ b/log_reasonings/2026-07-07-agent-fanout-strategy.md @@ -0,0 +1,44 @@ +# 2026-07-07 — Fanning out the production-RAG-hardening plan across agents + +## Context + +`docs/superpowers/plans/2026-07-07-production-rag-hardening.md` defines 15 +tasks. User asked to fan out one agent per task. Inspection of the plan's +per-task `**Files:**` lists showed heavy overlap: nearly every task touches +`localrag/settings.py` and `.env.example`; `localrag/api/service.py` is +touched by 8 of 15 tasks; `localrag/rag/retriever.py` and +`localrag/rag/engine.py` are each touched by 4-5 tasks. True 15-way parallel +execution against the same branch would guarantee merge conflicts. + +Explicit task dependencies also exist: Task 5 needs Task 2's metadata fields +(soft dependency — fields already existed conceptually, but sequencing is +safer); Task 6 needs Task 4's and Task 5's retriever shape; Task 15 needs +Task 4's `metadata_filter` explicitly; Task 14 needs Task 13's +`ResilientProvider` explicitly. + +## Decision + +Asked the user directly (via AskUserQuestion) rather than guessing: chose +**batched worktrees with sequential merge** — group tasks with zero file +overlap to run truly in parallel inside isolated `git worktree`s, merge each +completed task into `main` (rebasing onto `main` first if it moved since the +worktree branched), then launch the next batch. This trades some parallelism +for safety, appropriate given the file-contention density measured above. + +## Consequences + +- Every worktree branch needs a rebase-onto-`main` check before merge, since + `main` moves between when a worktree is created and when its agent + finishes — plain `git merge --ff-only` fails whenever another task landed + in between (happened for Task 3 and Task 4 in this run). +- Doc files (`docs/rag-retrieval.md` etc.) that multiple tasks append + sections to will conflict at rebase time even when the code changes don't + — these are trivial "keep both sections" resolutions, not real conflicts, + but still require a manual look rather than blind `--theirs`/`--ours`. +- A subagent hit the session's rate limit mid-task three times in this run + (not a code bug) — see `[[2026-07-07-agent-retry-discipline]]` for the + resulting persistence rule. + +## Related + +`[[2026-07-07-ragas-offline-judge]]`, `[[2026-07-07-agent-retry-discipline]]`. diff --git a/log_reasonings/2026-07-07-agent-retry-discipline.md b/log_reasonings/2026-07-07-agent-retry-discipline.md new file mode 100644 index 0000000..d68b173 --- /dev/null +++ b/log_reasonings/2026-07-07-agent-retry-discipline.md @@ -0,0 +1,46 @@ +# 2026-07-07 — Stop on repeated failure, don't grind + +## Context + +During the plan fan-out (see `[[2026-07-07-agent-fanout-strategy]]`), three +concurrent subagents failed mid-task with "You've hit your session limit" — +an infrastructure/rate-limit error, not a real code failure. Before that was +understood, agent transcripts showed repeated retry attempts against the same +failure. User's instruction: "if a test is failing repeatedly, be critical +instead of keep going on and on. Ask me when in doubts via a form." + +## Decision + +Adopted as a standing rule (saved to persistent memory, not just this log): +after 2-3 failed attempts at the same fix, stop and surface the failure +(to the user, or in a subagent's case, report back to the orchestrator) +rather than trying further variations blind. Distinguish: + +- **Real logic bugs** — worth some hypothesis-driven retry, per + `superpowers:systematic-debugging`. +- **Environment/infra errors** (session limits, network, unavailable + service) — report immediately, do not retry-loop at all; these will not + self-resolve by trying a different code fix. + +Applied concretely later the same session: when verifying the ragas judge +fix, `.env` in the worktree had no `OPENAI_API_KEY` (by design — `.env` is +gitignored, worktrees don't inherit it); rather than declaring the fix +"probably fine" after one failed local verification, copied the real key +over, re-checked, discovered the key itself was blank in the source `.env` +too, and escalated the real ambiguity (can this be verified locally at all?) +to the user via AskUserQuestion instead of asserting success on partial +evidence. + +## Consequences + +- Subagent prompts for this plan's remaining tasks now explicitly instruct: + "If you hit the same failure a third time with the same kind of fix, + STOP and report back the exact error instead of trying further + variations." +- Slower in the moment (asks more questions) but avoids burning tool-call + budget/context on unwinnable retries, and avoids false "verified working" + claims when verification was actually only partial. + +## Related + +`[[2026-07-07-agent-fanout-strategy]]`, `[[2026-07-07-ragas-offline-judge]]`. diff --git a/log_reasonings/2026-07-07-ragas-offline-judge.md b/log_reasonings/2026-07-07-ragas-offline-judge.md new file mode 100644 index 0000000..dec402e --- /dev/null +++ b/log_reasonings/2026-07-07-ragas-offline-judge.md @@ -0,0 +1,67 @@ +# 2026-07-07 — RAGAS judge LLM: stay offline, avoid new frameworks + +## Context + +Plan task "Gate RAGAS evals on PRs" hit a pre-existing break: `evals/run_evals.py` +imported legacy `ragas.metrics.collections` singletons; installed `ragas==0.4.3` +requires initialized `Metric` instances, and its per-metric result access +changed from scalar to per-row list. Separately, `ragas.evaluate()` defaults its +judge LLM to `OpenAI()`, which requires `OPENAI_API_KEY` — not present anywhere +in this project's `.env`, and fundamentally at odds with LocalRAG's +offline-first positioning ("Your documents, your models, your machine."). + +## Options considered + +1. **Keep OpenAI as judge**, rely on CI's `secrets.OPENAI_API_KEY`. Rejected — + contradicts offline-first positioning; can't be verified locally at all. +2. **LangChain wrapper** (`langchain-ollama` + `ragas.llms.LangchainLLMWrapper`). + Technically works, but user has a standing preference against pulling in + LangChain for this kind of glue — sees it as heavy/abstraction-for-its-own-sake + for what's fundamentally an HTTP call to a local model. +3. **Swap the whole eval framework to DeepEval** (suggested via pasted + third-party/Gemini advice, with a `prometheus2` judge model). Rejected as + disproportionate to the actual problem: this is a full eval-framework + migration (new dataset shape, new metric API, new dependency, new judge + model not even pulled yet), not a compatibility fix. DeepEval is not + obviously "lighter" than ragas — it's a different framework with its own + footprint, just not yet installed. Kept as a documented option for a future, + separately-scoped task if ever revisited. +4. **Ollama via LiteLLM adapter** (`ragas.llms.LiteLLMStructuredLLM`). Avoids + LangChain, but pulls in `litellm` + `instructor` for what amounts to talking + OpenAI-wire-protocol to a local server that already speaks that protocol. +5. **Ollama via its native OpenAI-compatible `/v1` endpoint**, using the + `openai` Python client already a core dependency of this project (LocalRAG + already supports `LLM_BACKEND=openai` as a pluggable backend). Ragas ships + `ragas.llms.llm_factory(...)` / `ragas.embeddings.OpenAIEmbeddings(...)` + that accept any OpenAI-shaped client — pointing it at + `http://localhost:11434/v1` with a dummy API key just works. + +## Decision + +Went with **option 5**. Zero new dependencies, no LangChain, no DeepEval +migration, no `prometheus2` download. `evals/run_evals.py` now builds an +`openai.OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")` client +and passes it into ragas's own `llm_factory`/`OpenAIEmbeddings` wrappers, judge +model defaulting to `gemma3:4b` (the project's own default LLM) and +`nomic-embed-text` for embeddings (the project's own default embedder) — both +already pulled by the existing Docker Ollama container. + +## Consequences + +- Evals are fully offline-capable and reproducible on any machine with the + project's own Docker stack running — no external API key needed anywhere, + including in CI (the `evals-pr.yml` gate no longer needs + `secrets.OPENAI_API_KEY`). +- Judge quality is bounded by a small local model (`gemma3:4b`, 4B params) + rather than GPT-4-class judges — thresholds in `PASS_THRESHOLDS` may need + recalibration once a full real run completes; flagged as a follow-up if + scores land far from expectations. +- `--judge-model`/`--ollama-url` CLI flags added so a different/larger local + judge can be swapped in without code changes. + +## Related + +See `[[2026-07-07-agent-fanout-strategy]]` for how this fix was executed as +part of a larger multi-agent plan rollout, and `docs/adr/` for LocalRAG's +existing formal architecture decisions (this file is a lighter-weight, +session-level reasoning log rather than a settled ADR). From 0da72ba591626a5886c4786832c011607a4b03c4 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 16:23:01 +0200 Subject: [PATCH 19/20] docs: add ADRs for resilient LLM routing, reranking, offline ragas judge Retroactively documents three architecture decisions from the production- RAG-hardening plan rollout that were missing formal ADRs (only captured in log_reasonings/ session notes): resilient/uniformly-routed LLM providers (Task 13-14), optional cross-encoder reranking (Task 6), and the offline Ollama-backed ragas judge (Task 1). Also clarifies in evals/run_evals.py that AsyncOpenAI/OpenAIEmbeddings there are OpenAI-shaped clients pointed at local Ollama, not a dependency on the OpenAI service. --- .../adr/007-resilient-llm-provider-routing.md | 42 ++++++++++++ docs/adr/008-cross-encoder-reranking.md | 44 +++++++++++++ docs/adr/009-offline-ragas-judge.md | 65 +++++++++++++++++++ evals/run_evals.py | 4 ++ 4 files changed, 155 insertions(+) create mode 100644 docs/adr/007-resilient-llm-provider-routing.md create mode 100644 docs/adr/008-cross-encoder-reranking.md create mode 100644 docs/adr/009-offline-ragas-judge.md diff --git a/docs/adr/007-resilient-llm-provider-routing.md b/docs/adr/007-resilient-llm-provider-routing.md new file mode 100644 index 0000000..fa77875 --- /dev/null +++ b/docs/adr/007-resilient-llm-provider-routing.md @@ -0,0 +1,42 @@ +# ADR 007: Resilient, uniformly-routed LLM provider abstraction + +## Context + +`LLM_BACKEND` (ollama/openai/anthropic) existed as a setting, but `RAGEngine` +hand-rolled its own httpx calls to Ollama directly for prompt generation and +streaming, bypassing `localrag/llm/factory.py::build_provider` entirely. +Changing `LLM_BACKEND` had no effect on `/query` — it only ever hit Ollama. +Separately, no provider had retry, timeout backoff, or fallback behavior: a +transient Ollama hiccup or OpenAI/Anthropic rate limit surfaced immediately +as a user-facing failure. + +## Decision + +- Give `BaseLLMProvider` two abstract methods, `generate_from_prompt` and + `stream_from_prompt`, implemented per backend (`ollama.py`, + `openai_provider.py`, `anthropic_provider.py`). +- Wrap every provider `build_provider` returns in `ResilientProvider` + (`localrag/llm/resilience.py`): tenacity retry with backoff, a pybreaker + circuit breaker, and an optional configured fallback backend + (`llm_fallback_backend`) that takes over once the breaker trips. +- Give `RAGEngine` a required `provider: BaseLLMProvider` field; route + `_stream_chat_tokens` through `self.provider.stream_from_prompt(...)` + instead of its own httpx client. `get_engine()` wires + `provider=build_provider(settings)`. + +## Consequences + +- `LLM_BACKEND` now genuinely governs `/query` — verified live by setting + `LLM_BACKEND=openai` with a bogus key and observing a real + `401 AuthenticationError` from the OpenAI SDK instead of a silent + fall-through to Ollama. +- All three backends get retry/circuit-breaker/fallback uniformly, instead of + resilience logic living only where someone happened to add it. +- Every LLM-calling code path (RAG answers, ragas eval judge, query + rewriting) now goes through one seam (`build_provider`), so future + resilience or observability changes apply everywhere at once. + +## Related + +`[[../architecture.md]]`, `docs/superpowers/plans/2026-07-07-production-rag-hardening.md` +(Tasks 13–14). diff --git a/docs/adr/008-cross-encoder-reranking.md b/docs/adr/008-cross-encoder-reranking.md new file mode 100644 index 0000000..d42a2d7 --- /dev/null +++ b/docs/adr/008-cross-encoder-reranking.md @@ -0,0 +1,44 @@ +# ADR 008: Optional cross-encoder reranking as a final relevance step + +## Context + +Hybrid vector+BM25 retrieval (ADR 005) with RRF fusion ranks well on lexical +and semantic signal separately, but neither embedding similarity nor BM25 +score directly optimizes for "is this chunk actually relevant to this exact +question" the way a cross-encoder (which scores the pair jointly) does. +Running a cross-encoder over the full corpus per query would be too slow; +running it only on the fused top-k throws away recall it could have fixed. + +## Decision + +Make reranking an optional, off-by-default final step: + +- When `RERANK_ENABLED=true`, `Retriever.retrieve` over-fetches + `RERANK_FETCH_K` candidates (instead of the usual `top_k * 2`) from the + vector/hybrid path. +- `CrossEncoderReranker` (`localrag/rag/reranker.py`), backed by a local + `sentence-transformers` cross-encoder model (`RERANK_MODEL`, default + `cross-encoder/ms-marco-MiniLM-L-6-v2`), scores each `(question, chunk)` + pair and trims back down to `top_k`. +- Reranking runs on the raw candidate list, strictly before freshness decay + (ADR 006) and parent-section expansion — it is the last relevance-ordering + step; freshness/expansion apply to the reranked, already-trimmed result. +- `sentence-transformers` is an optional dependency (`uv sync --extra + rerank`); nothing imports it unless the feature is enabled, keeping the + default install lightweight. + +## Consequences + +- Retrieval quality can improve for ambiguous queries without changing the + default (disabled) behavior or the default dependency footprint. +- Extra latency and a model download when enabled — appropriate to enable + per-deployment, not universally. +- `_fuse_results`'s dead RRF weight branch (only ever taken when + `bm25_weight == 0.5`, mathematically identical to the other branch) was + removed as part of this change since it was mechanically adjacent cleanup, + not a new decision in itself. + +## Related + +`[[005-hybrid-retrieval]]`, `[[006-freshness-decay]]`, +`docs/superpowers/plans/2026-07-07-production-rag-hardening.md` (Task 6). diff --git a/docs/adr/009-offline-ragas-judge.md b/docs/adr/009-offline-ragas-judge.md new file mode 100644 index 0000000..4d3b465 --- /dev/null +++ b/docs/adr/009-offline-ragas-judge.md @@ -0,0 +1,65 @@ +# ADR 009: RAGAS eval judge runs on local Ollama, not OpenAI + +## Context + +`evals/run_evals.py` gates the RAG pipeline's quality on faithfulness, +answer relevancy, context precision, and context recall (ragas metrics). +Ragas's `evaluate()` defaults its judge LLM to `OpenAI()`, which requires +`OPENAI_API_KEY` — not present anywhere in this project's `.env`, and at +odds with LocalRAG's offline-first positioning ("Your documents, your +models, your machine."). Installed `ragas==0.4.3` also deprecated the +`ragas.metrics.*` singleton import style this script originally used, in +favor of `ragas.metrics.collections.*` classes scored per-sample via async +`.ascore()`. + +## Decision + +Point the judge LLM and embeddings at the same local Ollama instance +LocalRAG itself uses, via Ollama's OpenAI-compatible `/v1` endpoint: + +```python +ollama_client = AsyncOpenAI(base_url=f"{ollama_url}/v1", api_key="ollama") +judge_llm = llm_factory(judge_model, client=ollama_client, adapter="instructor") +judge_embeddings = OpenAIEmbeddings(client=ollama_client, model="nomic-embed-text") +``` + +`openai.AsyncOpenAI`/`ragas.embeddings.OpenAIEmbeddings` are used here purely +as clients for the OpenAI-shaped wire protocol Ollama speaks — `base_url` +points at localhost, the API key is a required-but-ignored dummy string, and +no request ever reaches OpenAI's actual service. This was chosen over: + +1. **OpenAI as judge via CI secret** — rejected, can't be verified locally, + contradicts offline-first positioning. +2. **LangChain wrapper** (`langchain-ollama` + + `ragas.llms.LangchainLLMWrapper`) — works, but pulls in a dependency this + project's owner prefers to avoid for what's fundamentally an HTTP call to + a local model. +3. **Swap the whole eval framework to DeepEval** — rejected as a + disproportionate framework migration for a compatibility fix; DeepEval + isn't obviously lighter than ragas, just a different footprint. +4. **LiteLLM adapter** — avoids LangChain, but pulls in `litellm` + + `instructor` for the same job the `openai` client (already a core + dependency, since it backs the real `LLM_BACKEND=openai` provider) does + natively. + +Judge model defaults to `gemma3:4b`, embeddings to `nomic-embed-text` — the +project's own default LLM/embedder, already pulled by the existing Docker +Ollama container. Both are overridable via `--judge-model`/`--ollama-url`. + +## Consequences + +- Evals are fully offline-capable and reproducible on any machine running + the project's own Docker stack — no external API key needed anywhere, + including CI (`.github/workflows/evals-pr.yml` runs an `ollama/ollama` + service container and pulls the judge/embedding models before scoring). +- Judge quality is bounded by a small local model (`gemma3:4b`, 4B params) + rather than a GPT-4-class judge — `PASS_THRESHOLDS` may need recalibration + if scores drift far from expectations on a real corpus. First full run: + faithfulness 0.942, answer_relevancy 0.671, context_precision 0.891, + context_recall 1.000 — all above threshold. +- Zero new dependencies added for this fix. + +## Related + +`docs/superpowers/plans/2026-07-07-production-rag-hardening.md` (Task 1), +session narrative in `log_reasonings/2026-07-07-ragas-offline-judge.md`. diff --git a/evals/run_evals.py b/evals/run_evals.py index 1a7027a..26168fd 100644 --- a/evals/run_evals.py +++ b/evals/run_evals.py @@ -181,6 +181,10 @@ def main() -> None: rows = _build_rows(records, args.api_url, args.api_key, offline=args.offline) print(f"Running RAGAS evaluation (judge={args.judge_model} via {args.ollama_url})...") + # AsyncOpenAI/OpenAIEmbeddings here are just clients for the OpenAI-shaped + # wire protocol that Ollama's /v1 endpoint speaks. base_url points at the + # local Ollama instance, api_key is a required-but-ignored dummy value — + # no OpenAI account, key, or network call is ever involved. ollama_client = AsyncOpenAI(base_url=f"{args.ollama_url.rstrip('/')}/v1", api_key="ollama") judge_llm = llm_factory(args.judge_model, client=ollama_client, adapter="instructor") judge_embeddings = OpenAIEmbeddings(client=ollama_client, model="nomic-embed-text") From 93847a027eb413cbd469425dcc38ced5ceb65e58 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Wed, 8 Jul 2026 17:00:45 +0200 Subject: [PATCH 20/20] fix: remove private-method boundary leak, dedupe ingest retry loop RAGEngine._extract_sources was private but reached across the API/rag boundary via a noqa-suppressed access; made it public (extract_sources) instead. IngestionService.ingest_paths had two near-identical retry-loop bodies; extracted the shared per-file try/except into _safe_ingest_one, keeping exception-vs-no-chunks-produced distinguishable via a (result, error) tuple. --- localrag/api/repository.py | 2 +- localrag/api/service.py | 2 +- localrag/ingestion/service.py | 36 +++++++++++++++++++++++++---------- localrag/rag/engine.py | 4 ++-- tests/test_api.py | 2 +- tests/test_rag_engine.py | 2 +- 6 files changed, 32 insertions(+), 16 deletions(-) diff --git a/localrag/api/repository.py b/localrag/api/repository.py index 0238300..df7c177 100644 --- a/localrag/api/repository.py +++ b/localrag/api/repository.py @@ -12,7 +12,7 @@ class ChromaCollectionRepository: _vector_store: VectorStore def list_collection_names(self) -> list[str]: - return self._vector_store.list_collections() # now returns list[str] + return self._vector_store.list_collections() def delete_collection(self, name: str) -> None: self._vector_store.delete_collection(name) diff --git a/localrag/api/service.py b/localrag/api/service.py index 515ca61..29caf3e 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -356,7 +356,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) # noqa: SLF001 + raw_sources = engine.extract_sources(contexts) sources = [ SourceRef( source=str(s.get("source", "")), diff --git a/localrag/ingestion/service.py b/localrag/ingestion/service.py index 38ca50e..a8bfd29 100644 --- a/localrag/ingestion/service.py +++ b/localrag/ingestion/service.py @@ -108,16 +108,17 @@ def _stored_content_hashes(self) -> dict[str, str]: return hashes def ingest_paths(self, paths: list[Path], embed_model: str | None = None) -> IngestionResult: - total_chunks = 0 files_processed = 0 + total_chunks = 0 processed_sources: list[str] = [] retry_queue: list[Path] = [] for resolved_path in self._allowed_paths(paths): - try: - chunks_added = self._ingest_one(resolved_path, embed_model) - except Exception as exc: # retried once below, not fatal yet - logger.warning("ingest_file_failed_will_retry path=%s error=%s", resolved_path, exc) + 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: @@ -130,11 +131,12 @@ def ingest_paths(self, paths: list[Path], embed_model: str | None = None) -> Ing if retry_queue: logger.info("ingest_retry_start count=%s", len(retry_queue)) for resolved_path in retry_queue: - try: - chunks_added = self._ingest_one(resolved_path, embed_model) - except Exception as exc: # collected for the caller, not raised - logger.error("ingest_file_failed_permanently path=%s error=%s", resolved_path, exc) - failed_sources.append(FailedSource(source=str(resolved_path), error=str(exc))) + 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 @@ -153,6 +155,20 @@ def ingest_paths(self, paths: list[Path], embed_model: str | None = None) -> Ing failed_sources=failed_sources, ) + def _safe_ingest_one( + self, resolved_path: Path, embed_model: str | None + ) -> tuple[int | None, str | None]: + """Ingest one file, catching exceptions instead of raising. + + Returns `(chunks_added, None)` on success (`chunks_added` is None when the + file produced no chunks and was skipped), or `(None, error_message)` if + `_ingest_one` raised. Caller decides whether to log/retry/report the error. + """ + try: + return self._ingest_one(resolved_path, embed_model), None + except Exception as exc: + return None, str(exc) + def _allowed_paths(self, paths: list[Path]) -> list[Path]: allowed: list[Path] = [] for path in paths: diff --git a/localrag/rag/engine.py b/localrag/rag/engine.py index 552c2bc..5d516c4 100644 --- a/localrag/rag/engine.py +++ b/localrag/rag/engine.py @@ -102,10 +102,10 @@ def _stream_chat_tokens( if event["type"] == "token": yield event logger.info("rag_stream_done") - yield {"type": "final", "sources": self._extract_sources(contexts), "low_confidence": False} + 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, object]]: seen: set[tuple[str, int]] = set() sources: list[dict[str, object]] = [] for context in contexts: diff --git a/tests/test_api.py b/tests/test_api.py index d73a5a4..8ce8289 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -50,7 +50,7 @@ def stream_chat_from_contexts( yield {"type": "final", "sources": [{"source": "doc.md", "chunk_index": 1}]} @staticmethod - def _extract_sources(contexts: list[dict[str, Any]]) -> list[dict[str, object]]: + def extract_sources(contexts: list[dict[str, Any]]) -> list[dict[str, object]]: return [ {"source": str(c.get("source", "")), "chunk_index": int(c.get("chunk_index", -1))} for c in contexts diff --git a/tests/test_rag_engine.py b/tests/test_rag_engine.py index 714b58f..a215899 100644 --- a/tests/test_rag_engine.py +++ b/tests/test_rag_engine.py @@ -96,7 +96,7 @@ def test_rag_engine_extract_sources_includes_heading_path_and_chunk_type() -> No provider=FakeProvider(tokens=[]), ) - sources = engine._extract_sources(contexts) # noqa: SLF001 + sources = engine.extract_sources(contexts) assert sources == [ {