diff --git a/Dockerfile b/Dockerfile index 9f5b7b3..a75981f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,8 @@ FROM python:3.13-slim@sha256:6771159cd4fa5d9bba1258caf0b82e6b73458c694d178ad97c5e925c2d0e1a91 AS base +ARG LOCALRAG_BUILD_SHA=unknown +ENV LOCALRAG_BUILD_SHA=${LOCALRAG_BUILD_SHA} + WORKDIR /app # tesseract-ocr: required at runtime for scanned/image-only PDF pages (see docs/ocr.md). diff --git a/Taskfile.yml b/Taskfile.yml index 92d61d0..9a40ad1 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -127,7 +127,12 @@ tasks: docker-up: desc: Start the Compose stack in the background without removing volumes cmds: - - '{{.COMPOSE}} --project-name "{{.PROJECT}}" -f "{{.COMPOSE_FILE}}" {{if .COMPOSE_OVERRIDE}}-f "{{.COMPOSE_OVERRIDE}}" {{end}}up --build -d' + - 'LOCALRAG_BUILD_SHA="$$(git rev-parse HEAD)" {{.COMPOSE}} --project-name "{{.PROJECT}}" -f "{{.COMPOSE_FILE}}" {{if .COMPOSE_OVERRIDE}}-f "{{.COMPOSE_OVERRIDE}}" {{end}}up --build -d' + + docker-check: + desc: Compare the running API image identity with the current Git revision + cmds: + - '{{.PYTHON}} -c ''import json, os, subprocess, sys, urllib.request; expected=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(); request=urllib.request.Request("{{.API_URL}}/build-info", headers={"X-API-Key": os.environ["API_KEY"]}); actual=json.load(urllib.request.urlopen(request))["build_sha"]; print(f"running={actual} source={expected}"); sys.exit(0 if actual == expected else 1)''' docker-down: desc: Stop the Compose stack; named volumes are preserved diff --git a/docker-compose.yml b/docker-compose.yml index 7be9440..75fc4a9 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -31,7 +31,10 @@ services: # Runs from the same built image as localrag-api, so no separate script to # maintain. Exiting 0 is expected here — it's not meant to stay running. localrag-setup: - build: . + build: + context: . + args: + LOCALRAG_BUILD_SHA: ${LOCALRAG_BUILD_SHA:-unknown} image: localrag-api:local command: /app/.venv/bin/localrag setup environment: @@ -45,7 +48,10 @@ services: restart: "no" localrag-api: - build: . + build: + context: . + args: + LOCALRAG_BUILD_SHA: ${LOCALRAG_BUILD_SHA:-unknown} image: localrag-api:local restart: unless-stopped ports: @@ -81,7 +87,10 @@ services: start_period: 30s localrag-mcp: - build: . + build: + context: . + args: + LOCALRAG_BUILD_SHA: ${LOCALRAG_BUILD_SHA:-unknown} image: localrag-api:local restart: unless-stopped ports: diff --git a/docs/adr/040-request-scoped-collection-selection.md b/docs/adr/040-request-scoped-collection-selection.md new file mode 100644 index 0000000..16f7d0f --- /dev/null +++ b/docs/adr/040-request-scoped-collection-selection.md @@ -0,0 +1,34 @@ +# ADR 040: Request-Scoped Collection Selection + +## Status + +Accepted + +## Context + +The HTTP query API previously always used the collection configured at process +startup. That made `GET /collections` expose namespaces that HTTP clients could +not query, while the CLI could select them with a settings override. + +## Decision + +`POST /query`, `/query/contexts`, and `/query/stream` accept an optional +`collection` field. An omitted field retains the configured +`chroma_collection_name`. A supplied name selects an existing Chroma collection +for that request by creating a request-scoped vector store, BM25 snapshot, and +retriever/engine view. Process-wide settings and the cached default engine are +not mutated. + +The endpoint remains behind `API_KEY`. Collection names are not treated as +tenant authorization: deployments that use collections for tenant separation +must enforce per-key authorization at a gateway or provide separate persist +paths. This is an explicit single-deployment namespace selection feature, not a +multi-tenant isolation boundary. + +## Consequences + +- JSON, contexts, and SSE query paths can address any existing collection. +- A request-scoped hybrid query rebuilds the BM25 snapshot for the selected collection. +- Unknown collections return a query `404` rather than exposing a backend error. +- Retriever plugins that do not expose the built-in collection seam cannot use + request-scoped selection. diff --git a/docs/adr/README.md b/docs/adr/README.md index 579a99c..a4e4f17 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -61,6 +61,7 @@ default. | [037](037-grouped-configuration-model.md) | Grouped configuration model behind flat public names | Accepted | Configuration | | [038](038-application-and-mcp-boundaries.md) | Transport-agnostic application boundary and MCP adapter | Amended by [039](039-fastmcp-sdk-adoption.md) | Architecture / MCP | | [039](039-fastmcp-sdk-adoption.md) | Adopt the FastMCP SDK for the MCP adapter | Accepted | Architecture / MCP | +| [040](040-request-scoped-collection-selection.md) | Request-scoped HTTP collection selection | Accepted | API / Retrieval | ## Research spikes diff --git a/docs/agent-friction.md b/docs/agent-friction.md new file mode 100644 index 0000000..c6d2b2f --- /dev/null +++ b/docs/agent-friction.md @@ -0,0 +1,37 @@ +# Agent Friction Notes + +This is a short record of repository traps discovered while resolving and +verifying changes. Keep entries concrete and delete them only when the +underlying workflow or contract changes. + +## Merge Conflicts + +PR #169 conflicted with the progress-output work because both branches changed +the same CLI ingest section, ingest entry point, and ingestion-service tests. +The correct resolution was additive: keep progress callbacks and wrap the whole +write path in the cross-process ingest lock. Rebase the PR branch onto current +`main` before reconstructing tests; resolving only the documentation conflict +leaves the test file structurally broken. + +## Docker Drift + +An already-running Compose stack serves the old image even when the source tree +has changed. Rebuild before integration tests. The image now carries +`LOCALRAG_BUILD_SHA`; `task docker-up` stamps the current revision and +`task docker-check` compares it through the authenticated `/build-info` +endpoint. + +## Chroma Lifecycle + +Deleting a collection removes Chroma metadata but can leave its persisted HNSW +directory. Resolve the vector segment ID before deletion, remove only that +directory after successful metadata deletion, and invalidate cached retrieval +objects. Otherwise the next query in the same long-running API can use the +deleted collection UUID and return a misleading 503. + +## Integration Test Assumptions + +The Compose setup pulls `gemma3:4b`; integration tests must not request an +unprovisioned model such as `qwen2.5:0.5b`. Run the full integration suite only +after rebuilding the image and starting the stack, using +`LOCALRAG_TEST_API_KEY` for protected endpoints. diff --git a/docs/agent-navigation.md b/docs/agent-navigation.md index b53eaaa..267d0d6 100644 --- a/docs/agent-navigation.md +++ b/docs/agent-navigation.md @@ -36,7 +36,7 @@ override unless `COMPOSE_OVERRIDE` is supplied. | Transport-agnostic use cases | `localrag/application/`, `localrag/application/service.py`, `localrag/application/container.py` | | MCP tools and transports | `localrag/mcp/server.py`, `localrag/mcp/app.py`, `docs/mcp.md` | | API request/response OpenAPI models | `localrag/api/schemas.py` | -| Application use cases (health, ingest rules, query JSON + SSE, collections including rebuild) | `localrag/application/service.py` | +| Application use cases (health, ingest rules, query JSON + SSE, request-scoped collections, collections including rebuild) | `localrag/application/service.py`, `localrag/rag/engine.py`, `localrag/rag/retriever.py` | | API schema/error adapter | `localrag/api/service.py`, `localrag/api/exceptions.py` | | Application persistence boundary (Chroma collections) | `localrag/application/repository.py` | | API app factory (lifespan, middleware, error handlers) | `localrag/api/main.py` | @@ -48,7 +48,7 @@ override unless `COMPOSE_OVERRIDE` is supplied. | Log format, levels, request ID | `localrag/logging_config.py`, `localrag/api/middleware.py`, `LOG_LEVEL` in `localrag/settings.py` | | Optional tracing / observability | `localrag/observability/tracing.py`, `OTEL_*` in `localrag/settings.py`, [observability.md](observability.md), [ADR 030](adr/030-optional-otel-observability-boundary.md) | | API key auth | `localrag/api/dependencies.py` (`require_api_key`), `API_KEY` in `localrag/settings.py` | -| Prometheus metrics endpoint | `localrag/api/routers/metrics.py` | +| Prometheus metrics and running build identity | `localrag/api/routers/metrics.py`, `Taskfile.yml`, `Dockerfile`, [observability.md](observability.md) | | LLM provider abstraction | `localrag/llm/providers/`, `localrag/llm/factory.py` | | Cost estimation | `localrag/llm/costs.py` | | Agent tool-use (search_documents / answer_directly) | `localrag/agent/service.py`, `localrag/api/routers/agent.py` | @@ -63,6 +63,7 @@ override unless `COMPOSE_OVERRIDE` is supplied. | Ingestion embedding cache | `localrag/embedding/cache.py`, embedding cache settings, [ADR 024](adr/024-embedding-cache-contract.md), `benchmarks/embedding_cache_benchmark.py` | | Ingest orchestration | `localrag/ingestion/service.py` | | Chroma collection / persist path | `localrag/storage/vector_store.py`, settings | +| Concurrent-ingest exclusion (one writer per persist path) | `localrag/storage/persist_lock.py`, `localrag/ingestion/service.py`, [ADR 035](adr/035-atomic-ingestion-replacement.md), [cli.md](cli.md) | | Retrieval mode / hybrid ranking / freshness decay / HyDE experiment | `localrag/rag/retriever.py`, `localrag/rag/hyde.py`, `localrag/rag/bm25_index.py`, `localrag/settings.py`, [ADR 025](adr/025-hyde-retrieval-experiment.md) | | Retriever plugin contract / discovery | `localrag/plugins/retriever.py`, [plugin-author-guide.md](plugin-author-guide.md), [ADR 032](adr/032-retriever-plugin-contract.md) | | Bounded adaptive retrieval policy / trace | `localrag/rag/adaptive.py`, `localrag/rag/engine.py`, adaptive settings, [ADR 023](adr/023-bounded-adaptive-retrieval.md) | @@ -102,6 +103,8 @@ uv run ruff check . Pre-commit and contribution workflow: [`CONTRIBUTING.md`](../CONTRIBUTING.md). For portable task commands and their variable contract, use [`Taskfile.yml`](../Taskfile.yml). +Known merge, Docker, Chroma, and integration-test friction is recorded in +[`agent-friction.md`](agent-friction.md). ## External dependencies diff --git a/docs/architecture.md b/docs/architecture.md index d7c1d86..e5089fc 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -62,10 +62,11 @@ flowchart LR ``` - **Ingest:** files → `loader` / `ingestion/parsers/*` → text → the shared `Chunk` contract (`localrag/ingestion/contract.py`) implemented by fixed, structural, or recursive strategies → the factory-created `EmbeddingProvider` → `VectorStore` (Chroma, persistent path from settings). Contract IDs are deterministic from source, strategy, index, and text; offsets are explicitly absent because current strategies normalize or repack text. Empty input emits no chunks and oversized atomic input is retained with an `oversized` marker. The same provider instance embeds retrieval queries. Collection metadata records provider/model/dimension and rejects incompatible operations; changing the embedding space requires an explicit rebuild. The application ingest use cases in `localrag/application/service.py` own path decode, existence checks, `INGEST_ROOTS`, upload limits, and background jobs; the HTTP adapter maps their DTOs and errors to OpenAPI responses. CLI and MCP calls use the same application use cases directly. See [ADR 021](adr/021-chunking-strategy-contract.md) and [ADR 038](adr/038-application-and-mcp-boundaries.md). + - **One writer per persist path:** Chroma's embedded client holds HNSW segments in per-process memory with no cross-process invalidation, so concurrent writers from separate processes silently lose writes ([ADR 035](adr/035-atomic-ingestion-replacement.md) already puts them out of contract). `IngestionService.ingest_paths` and `rebuild_collection` therefore take `localrag/storage/persist_lock.py::ingest_lock` — an advisory `flock` on `/.ingest.lock` — for the duration of the write; a competing process gets `ConcurrentIngestError` immediately (CLI: stderr message and exit `1`; HTTP: `409 Conflict`). The lock deliberately sits at the ingest use-case boundary rather than inside `VectorStore`, because `api/dependencies.py` and `application/runtime.py` cache a `VectorStore` for the whole process lifetime and a store-scoped lock would let a running API block every CLI ingest forever. Read and query paths are never locked, and an unopenable persist directory or a filesystem without `flock` degrades to a logged warning rather than a hard failure. `VectorStore.delete_collection` removes the deleted collection's recorded persisted HNSW segment directory after Chroma metadata deletion, without sweeping unrelated directories. - **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 the application `query_json` use case, adapted by `localrag/api/service.py`. Retrieval supports vector-only and hybrid (vector + BM25 with reciprocal-rank fusion), optional bounded query expansion, then applies optional freshness decay based on chunk `ingested_at`. When `ADAPTIVE_ENABLED=true`, `AdaptiveRetrievalPolicy` performs bounded evidence evaluation/escalation/refinement and adds an observable trace; thresholds are corpus-tuned heuristics, not calibrated confidence. JSON and SSE use the same engine policy path. `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 (JSON):** `POST /query` returns a complete `QueryResponse` (answer, sources, latency_ms, model) from the application `query_json` use case, adapted by `localrag/api/service.py`. Retrieval supports vector-only and hybrid (vector + BM25 with reciprocal-rank fusion), optional bounded query expansion, then applies optional freshness decay based on chunk `ingested_at`. When `ADAPTIVE_ENABLED=true`, `AdaptiveRetrievalPolicy` performs bounded evidence evaluation/escalation/refinement and adds an observable trace; thresholds are corpus-tuned heuristics, not calibrated confidence. JSON and SSE use the same engine policy path. `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. An optional request `collection` selects another collection for that request only; omitted values use `chroma_collection_name` ([ADR 040](adr/040-request-scoped-collection-selection.md)). 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. +- **Metrics:** `GET /metrics` exposes Prometheus metrics via `prometheus_client` (router at `localrag/api/routers/metrics.py`). Metrics and `GET /build-info` require `X-API-Key` when `API_KEY` is set; build info reports `LOCALRAG_BUILD_SHA` for `task docker-check`. ## Package map @@ -89,6 +90,7 @@ flowchart LR | Embedding | `localrag/embedding/`, `localrag/ingestion/embedder.py` | Provider protocol/factory, Ollama **`POST /api/embed`**, optional sentence-transformers backend, collection identity checks, and the disabled-by-default provider-aware ingestion vector cache | | Embedding cache | `localrag/embedding/cache.py` | Versioned hashed vector-only disk entries, atomic writes, process/thread locking, checksum validation, bounded LRU cleanup, and fail-open cache I/O | | Storage | `localrag/storage/vector_store.py` | Chroma client wrapper | +| Ingest ownership lock | `localrag/storage/persist_lock.py` | `ingest_lock` — advisory `flock` on `/.ingest.lock` giving one process exclusive write ownership of a Chroma persist path; fails fast with `ConcurrentIngestError` | | RAG | `localrag/rag/retriever.py`, `bm25_index.py`, `engine.py`, `prompt.py` | Hybrid retrieval (vector + BM25), freshness decay reranking, prompt build, LLM call | | Retriever plugins | `localrag/plugins/retriever.py`, `docs/plugin-author-guide.md` | Versioned `localrag.retrievers` entry points; deterministic selection and lifecycle ownership | | Context compression | `localrag/rag/compressor.py` | Disabled-by-default deterministic extractive compression after parent expansion; preserves retrieval provenance and hard token/character budgets | diff --git a/docs/cli.md b/docs/cli.md index 413a714..2d3a63a 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -58,6 +58,36 @@ Files that no parser can handle — and binary content in a file with a textual extension — are reported and skipped rather than being read as text. See [document-formats.md](document-formats.md). +`ingest` accepts a file or a directory and writes into the configured Chroma +persist path: + +```bash +uv run localrag ingest ./docs +uv run localrag ingest ./docs/guide.md +``` + +**One writer per persist path.** Chroma's embedded client keeps its HNSW +segments in per-process memory with no cross-process invalidation, so two +processes writing the same `CHROMA_PERSIST_PATH` corrupt each other's view — +the boundary [ADR 035](adr/035-atomic-ingestion-replacement.md) already declares +out of contract. Every ingest therefore takes an exclusive `flock` on +`/.ingest.lock` for the duration of the run, and the same +lock covers collection rebuilds. + +A second concurrent ingest **fails immediately** rather than queueing — an +ingest can run for minutes, so a caller is better told to retry than left +hanging: + +``` +status=error reason=concurrent_ingest detail=another ingest is already running against ./data/chroma; wait for it to finish or use a different collection +``` + +The message goes to stderr and the command exits `1`. The equivalent HTTP +ingest endpoints return `409 Conflict` with the same detail. Query and other +read paths are never locked. The lock is advisory: if the persist directory +cannot be opened, or the filesystem does not support `flock` (some network +mounts) the ingest proceeds unguarded and logs a warning. + ## Inspect `inspect` is read-only and never creates a missing collection, calls an LLM, or diff --git a/docs/deployment.md b/docs/deployment.md index 2093fd5..c21cb34 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -16,6 +16,11 @@ Set secrets in the environment or `.env` before starting: API_KEY='change-me' GRAFANA_ADMIN_PASSWORD='change-me-too' docker compose up -d ``` +`task docker-up` passes the current Git revision into the image as +`LOCALRAG_BUILD_SHA`. Run `API_KEY=... task docker-check` after source changes; +it calls the protected `/build-info` endpoint and exits non-zero when the +running image is stale. Rebuild the stack before trusting end-to-end results. + Ollama runs without a default GPU reservation so the stack also works on CPU CI runners. Configure GPU scheduling separately for hosts that support it. diff --git a/docs/observability.md b/docs/observability.md index e0c6184..0ec3736 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -35,7 +35,7 @@ telemetry. ## Prometheus Metrics -The `/metrics` endpoint exposes bounded counters and histograms for query +The authenticated `/metrics` endpoint exposes bounded counters and histograms for query duration, retrieved chunks, generated tokens, query/provider failures, cache hits and misses, ingestion failures, background job terminal status, upload cleanup, audit-log rotation, and HTTP failures by status class. Labels are limited to transport/outcome @@ -43,7 +43,12 @@ cleanup, audit-log rotation, and HTTP failures by status class. Labels are limit paths, and request IDs are never labels. JSON, adaptive, and SSE queries record the same duration, retrieval, token, audit, and failure signals. Fallback use, upload quota rejection, oversized audit records, and audit write/cleanup - failures have dedicated bounded counters. +failures have dedicated bounded counters. + +`GET /build-info` is protected by the same API key and returns the image's +`LOCALRAG_BUILD_SHA`. `task docker-up` stamps this value from the current Git +revision, and `task docker-check` compares it with the running API so a stale +Compose image is visible before diagnosing application behavior. Useful alerts include `rate(localrag_query_failures_total[5m]) > 0`, a sustained `localrag_ingest_jobs_pending` near the configured cap, and any increase in diff --git a/localrag/api/dependencies.py b/localrag/api/dependencies.py index ab9227c..a08ef56 100644 --- a/localrag/api/dependencies.py +++ b/localrag/api/dependencies.py @@ -8,6 +8,7 @@ from localrag.application.jobs import JobRegistry from localrag.application.repository import ChromaCollectionRepository +from localrag.application.runtime import clear_runtime_caches from localrag.embedding.base import EmbeddingProvider from localrag.embedding.cache import EmbeddingCache from localrag.embedding.factory import build_embedding_provider @@ -122,3 +123,15 @@ def get_collection_repository( store: VectorStore = Depends(get_vector_store), ) -> ChromaCollectionRepository: return ChromaCollectionRepository(_vector_store=store) + + +def invalidate_retrieval_caches() -> None: + """Drop collection-bound retrieval objects after a collection mutation.""" + if get_retriever.cache_info().currsize: + close = getattr(get_retriever(), "close", None) + if close is not None: + close() + get_engine.cache_clear() + get_retriever.cache_clear() + get_bm25_index.cache_clear() + clear_runtime_caches() diff --git a/localrag/api/exceptions.py b/localrag/api/exceptions.py index afcb43b..c7e8f46 100644 --- a/localrag/api/exceptions.py +++ b/localrag/api/exceptions.py @@ -3,6 +3,7 @@ from fastapi import status from localrag.application.errors import ApplicationError, ApplicationErrorKind +from localrag.storage.persist_lock import ConcurrentIngestError class HttpMappedError(Exception): @@ -26,6 +27,14 @@ class AgentApiError(HttpMappedError): """Raised when the agent endpoint cannot run (e.g. missing provider credentials).""" +class ConcurrentIngestApiError(HttpMappedError): + """Raised when another process already owns the Chroma persist path (ADR 035).""" + + @classmethod + def create(cls, exc: ConcurrentIngestError) -> ConcurrentIngestApiError: + return cls(status.HTTP_409_CONFLICT, str(exc)) + + def to_http_error(exc: ApplicationError) -> HttpMappedError: status_by_kind = { ApplicationErrorKind.BAD_REQUEST: status.HTTP_400_BAD_REQUEST, diff --git a/localrag/api/main.py b/localrag/api/main.py index 09ccf1f..87d303a 100644 --- a/localrag/api/main.py +++ b/localrag/api/main.py @@ -10,7 +10,7 @@ from fastapi.responses import JSONResponse from localrag.api.dependencies import get_embedder, get_retriever -from localrag.api.exceptions import HttpMappedError, to_http_error +from localrag.api.exceptions import ConcurrentIngestApiError, HttpMappedError, to_http_error from localrag.api.middleware import RequestContextMiddleware from localrag.api.routers.agent import router as agent_router from localrag.api.routers.collections import router as collections_router @@ -22,6 +22,7 @@ from localrag.logging_config import configure_logging from localrag.observability.tracing import configure_tracing, shutdown_tracing from localrag.settings import get_settings, load_settings, set_current_settings +from localrag.storage.persist_lock import ConcurrentIngestError logger = logging.getLogger(__name__) @@ -74,6 +75,13 @@ async def application_error_handler(request: Request, exc: ApplicationError) -> return await http_mapped_error_handler(request, to_http_error(exc)) +@app.exception_handler(ConcurrentIngestError) +async def concurrent_ingest_error_handler( + request: Request, exc: ConcurrentIngestError +) -> JSONResponse: + return await http_mapped_error_handler(request, ConcurrentIngestApiError.create(exc)) + + @app.exception_handler(Exception) async def unhandled_exception_handler(request: Request, exc: Exception) -> JSONResponse: logger.error( diff --git a/localrag/api/routers/collections.py b/localrag/api/routers/collections.py index ca0c76d..b6f2ee0 100644 --- a/localrag/api/routers/collections.py +++ b/localrag/api/routers/collections.py @@ -7,6 +7,7 @@ get_collection_repository, get_ingestion_service, get_query_cache, + invalidate_retrieval_caches, require_api_key, ) from localrag.api.schemas import ( @@ -42,6 +43,7 @@ def delete_collection( ) -> CollectionDeleteResponse: response = api_service.delete_collection_response(collection_repo, name) query_cache.clear() + invalidate_retrieval_caches() return response diff --git a/localrag/api/routers/metrics.py b/localrag/api/routers/metrics.py index d3400e4..5dd1a7e 100644 --- a/localrag/api/routers/metrics.py +++ b/localrag/api/routers/metrics.py @@ -1,13 +1,23 @@ from __future__ import annotations -from fastapi import APIRouter +import os + +from fastapi import APIRouter, Depends from fastapi.responses import Response from prometheus_client import CONTENT_TYPE_LATEST, generate_latest -router = APIRouter(prefix="", tags=["metrics"]) +from localrag.api.dependencies import require_api_key + +router = APIRouter(prefix="", tags=["metrics"], dependencies=[Depends(require_api_key)]) @router.get("/metrics", summary="Prometheus metrics") def metrics() -> Response: """Expose Prometheus metrics in text format for scraping.""" return Response(content=generate_latest(), media_type=CONTENT_TYPE_LATEST) + + +@router.get("/build-info", summary="Build identity") +def build_info() -> dict[str, str]: + """Expose the image identity so a running stack can be compared with source.""" + return {"build_sha": os.environ.get("LOCALRAG_BUILD_SHA", "unknown")} diff --git a/localrag/api/schemas.py b/localrag/api/schemas.py index 21daa1d..08318e5 100644 --- a/localrag/api/schemas.py +++ b/localrag/api/schemas.py @@ -68,6 +68,14 @@ class QueryRequest(BaseModel): ), examples=[{"source": "/docs/handbook.pdf"}], ) + collection: str | None = Field( + default=None, + description=( + "Optional Chroma collection to query. If omitted, uses the server's " + "configured `chroma_collection_name`." + ), + examples=["experiments"], + ) class IngestFileRequest(BaseModel): diff --git a/localrag/api/service.py b/localrag/api/service.py index 5c8d346..4935ac9 100644 --- a/localrag/api/service.py +++ b/localrag/api/service.py @@ -147,6 +147,7 @@ def query_json( model=request.model, n_results=request.n_results, metadata_filter=request.metadata_filter, + collection=request.collection, ), engine, query_cache, @@ -169,6 +170,7 @@ def get_query_contexts(request: schemas.QueryRequest, engine: RAGEngine) -> list model=request.model, n_results=request.n_results, metadata_filter=request.metadata_filter, + collection=request.collection, ), engine, ) @@ -185,6 +187,7 @@ def iter_query_sse_events( model=request.model, n_results=request.n_results, metadata_filter=request.metadata_filter, + collection=request.collection, ), engine, contexts, diff --git a/localrag/application/dto.py b/localrag/application/dto.py index bbb711e..b21e5e6 100644 --- a/localrag/application/dto.py +++ b/localrag/application/dto.py @@ -10,6 +10,7 @@ class QueryRequest: model: str | None = None n_results: int | None = None metadata_filter: dict[str, str] | None = None + collection: str | None = None @dataclass(frozen=True) diff --git a/localrag/application/service.py b/localrag/application/service.py index 773d7f6..08e3ba1 100644 --- a/localrag/application/service.py +++ b/localrag/application/service.py @@ -12,6 +12,7 @@ from uuid import uuid4 import httpx +from chromadb.errors import NotFoundError from localrag import metrics as app_metrics from localrag.application.dto import ( @@ -415,6 +416,7 @@ def query_json( # noqa: C901, PLR0915 request: QueryRequest, engine: RAGEngine, query_cache: QueryCache | None = None ) -> QueryResponse: """Blocking JSON query — retrieves context then generates a full answer.""" + engine = _engine_for_request(request, engine) t0 = time.perf_counter() cache_key: str | None = None if query_cache is not None: @@ -575,6 +577,7 @@ def _query_error_kind(exc: RetrievalError) -> ApplicationErrorKind: def get_query_contexts(request: QueryRequest, engine: RAGEngine) -> list[dict[str, Any]]: """Retrieve chunks synchronously so embedding / vector errors map to HTTP before SSE starts.""" + engine = _engine_for_request(request, engine) try: with span(SpanName.RETRIEVAL, {"stage": "retrieve"}): return engine.retriever.retrieve( @@ -592,6 +595,7 @@ def iter_query_sse_events( engine: RAGEngine, contexts: list[dict[str, Any]], ) -> Iterator[dict[str, Any]]: + engine = _engine_for_request(request, engine) t0 = time.perf_counter() logger.info( "query_start model=%s n_results=%s question_chars=%s", @@ -647,3 +651,15 @@ def iter_query_sse_events( "event": "error", "data": json.dumps({"detail": "LLM provider request failed."}), } + + +def _engine_for_request(request: QueryRequest, engine: RAGEngine) -> RAGEngine: + if request.collection is None or request.collection == engine.settings.chroma_collection_name: + return engine + try: + return engine.for_collection(request.collection) + except NotFoundError as exc: + raise QueryError( + ApplicationErrorKind.NOT_FOUND, + f"Collection '{request.collection}' not found.", + ) from exc diff --git a/localrag/cli/commands/ingest.py b/localrag/cli/commands/ingest.py index 32f3bb5..4ff3eaf 100644 --- a/localrag/cli/commands/ingest.py +++ b/localrag/cli/commands/ingest.py @@ -7,6 +7,7 @@ from localrag.application.container import get_ingestion_service from localrag.ingestion.service import IngestProgress +from localrag.storage.persist_lock import ConcurrentIngestError logger = logging.getLogger(__name__) @@ -32,10 +33,17 @@ def ingest(path: str, recursive: bool | None = None, *, quiet: bool = False) -> # run is indistinguishable from a hung one. on_progress = None if quiet else _echo_progress - if target.is_dir(): - result = service.ingest_directory(path=target, recursive=recursive, on_progress=on_progress) - else: - result = service.ingest_file(path=target, on_progress=on_progress) + try: + if target.is_dir(): + result = service.ingest_directory( + path=target, recursive=recursive, on_progress=on_progress + ) + else: + result = service.ingest_file(path=target, on_progress=on_progress) + except ConcurrentIngestError as exc: + logger.error("cli_ingest_conflict error=%s", exc) + typer.echo(f"status=error reason=concurrent_ingest detail={exc}", err=True) + raise typer.Exit(code=1) from exc failed_count = len(result.failed_sources) logger.info( diff --git a/localrag/ingestion/service.py b/localrag/ingestion/service.py index 7a359ed..c68ed5c 100644 --- a/localrag/ingestion/service.py +++ b/localrag/ingestion/service.py @@ -21,6 +21,7 @@ from localrag.observability.tracing import SpanName, span from localrag.rag.bm25_index import Bm25Index from localrag.settings import Settings, is_path_allowed +from localrag.storage.persist_lock import ingest_lock from localrag.storage.vector_store import VectorStore logger = logging.getLogger(__name__) @@ -108,7 +109,9 @@ def ingest_directory( return self.ingest_paths(files, embed_model=embed_model, on_progress=on_progress) def rebuild_collection(self, embed_model: str | None = None) -> RebuildCollectionResult: - with self._write_lock: + # Rebuild deletes and re-embeds every source, so it needs the same cross-process + # ownership of the persist path that a plain ingest does (ADR 035). + with ingest_lock(self.settings.chroma_persist_path), self._write_lock: sources = self.vector_store.list_distinct_sources() stored_hashes = self._stored_content_hashes() missing_sources: list[str] = [] @@ -152,7 +155,9 @@ def ingest_paths( embed_model: str | None = None, on_progress: ProgressCallback | None = None, ) -> IngestionResult: - with self._write_lock: + # The file lock coordinates separate processes; the RLock handles writers + # within this process and remains reentrant for rebuild delegation. + with ingest_lock(self.settings.chroma_persist_path), self._write_lock: return self._ingest_paths_locked(paths, embed_model, on_progress=on_progress) def _ingest_paths_locked( diff --git a/localrag/plugins/retriever.py b/localrag/plugins/retriever.py index ded1f89..1e259a0 100644 --- a/localrag/plugins/retriever.py +++ b/localrag/plugins/retriever.py @@ -157,9 +157,12 @@ def create(self, plugin_id: str, settings: Settings) -> RetrieverContract: raise PluginExecutionError( _message("Unable to create retriever plugin '{}': {}", plugin_id, exc) ) from exc + self.track(instance, plugin_id) + return instance + + def track(self, instance: RetrieverContract, plugin_id: str) -> None: self._instances.append(instance) self._instance_ids[id(instance)] = plugin_id - return instance def retrieve( self, @@ -203,6 +206,14 @@ def retrieve( ) -> list[RetrievalContext]: return self._registry.retrieve(self._instance, question, n_results, metadata_filter) + def for_collection(self, collection: str) -> ManagedRetriever: + """Create a tracked built-in retriever for a request-scoped collection.""" + if not isinstance(self._instance, Retriever): + raise TypeError("Per-request collections require the built-in retriever.") + instance = self._instance.for_collection(collection) + self._registry.track(instance, "builtin") + return type(self)(self._registry, instance) + def close(self) -> None: self._registry.close() diff --git a/localrag/rag/engine.py b/localrag/rag/engine.py index 5e6d442..314e10b 100644 --- a/localrag/rag/engine.py +++ b/localrag/rag/engine.py @@ -22,6 +22,16 @@ class RAGEngine: retriever: Retriever provider: BaseLLMProvider + def for_collection(self, collection: str) -> RAGEngine: + """Create a request-scoped engine targeting one collection.""" + if not hasattr(self.retriever, "for_collection"): + raise TypeError("Per-request collections require the built-in retriever.") + return type(self)( + settings=self.settings.with_overrides(chroma_collection_name=collection), + retriever=self.retriever.for_collection(collection), + provider=self.provider, + ) + def answer( self, question: str, diff --git a/localrag/rag/retriever.py b/localrag/rag/retriever.py index 14281bb..7df8754 100644 --- a/localrag/rag/retriever.py +++ b/localrag/rag/retriever.py @@ -66,6 +66,17 @@ class Retriever: reranker: CrossEncoderReranker | None = None last_hyde: HydeObservation | None = None + def for_collection(self, collection: str) -> Retriever: + """Create a request-scoped retriever for another collection.""" + store = self.vector_store.for_collection(collection) + return type(self)( + settings=self.settings.with_overrides(chroma_collection_name=collection), + embedder=self.embedder, + vector_store=store, + bm25_index=Bm25Index.from_vector_store(store) if self.bm25_index is not None else None, + reranker=self.reranker, + ) + def retrieve( self, question: str, diff --git a/localrag/storage/persist_lock.py b/localrag/storage/persist_lock.py new file mode 100644 index 0000000..5722e05 --- /dev/null +++ b/localrag/storage/persist_lock.py @@ -0,0 +1,118 @@ +"""Cross-process exclusion for writers against one Chroma persist directory. + +Chroma's ``PersistentClient`` keeps its HNSW segments in per-process memory with no +cross-process invalidation, so two processes writing the same persist path corrupt +each other's view (see ADR 035, which already declares multi-process writers out of +contract). This module is that external ownership boundary: an advisory ``flock`` on +``/.ingest.lock``. + +It deliberately lives outside ``VectorStore``: the API holds an ``lru_cache``d store +for the whole process lifetime, so a store-scoped lock would block every CLI ingest +forever. Acquire it around an ingest instead, and leave read paths unlocked. +""" + +from __future__ import annotations + +import logging +import threading +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import TextIO + +try: + import fcntl +except ImportError: # pragma: no cover - LocalRAG currently supports Unix hosts. + fcntl = None # type: ignore[assignment] + +logger = logging.getLogger(__name__) + +LOCK_FILE_NAME = ".ingest.lock" + + +class ConcurrentIngestError(RuntimeError): + """Raised when another process already owns the ingest lock for a persist path.""" + + +class _ReentrantFileLock: + """Holds one ``flock`` per persist path and counts nested acquisitions in-process. + + ``flock`` is per-file-description, so a second acquisition inside the same process + would silently succeed and its release would drop the lock while the outer scope + still believed it held it. Counting keeps nesting (ingest inside rebuild) correct. + """ + + def __init__(self) -> None: + self._guard = threading.Lock() + self._depth: dict[Path, int] = {} + self._files: dict[Path, TextIO] = {} + + def acquire(self, lock_path: Path) -> bool: + """Take the lock, returning whether a release is owed. Raises only on contention.""" + with self._guard: + if self._depth.get(lock_path, 0) > 0: + self._depth[lock_path] += 1 + return True + try: + lock_path.parent.mkdir(parents=True, exist_ok=True) + descriptor = lock_path.open("a+", encoding="utf-8") + except OSError: + # An unwritable persist directory cannot be ingested into anyway; let the + # store raise the real error instead of masking it as a lock failure. + logger.warning("ingest_lock_unavailable path=%s", lock_path) + return False + try: + if fcntl is not None: + fcntl.flock(descriptor.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + descriptor.close() + message = ( + f"another ingest is already running against {lock_path.parent}; " + "wait for it to finish or use a different collection" + ) + raise ConcurrentIngestError(message) from exc + except OSError: + # Filesystems without flock support (some network mounts) still ingest. + logger.warning("ingest_lock_unsupported path=%s", lock_path) + descriptor.close() + return False + self._files[lock_path] = descriptor + self._depth[lock_path] = 1 + return True + + def release(self, lock_path: Path) -> None: + with self._guard: + remaining = self._depth.get(lock_path, 0) - 1 + if remaining > 0: + self._depth[lock_path] = remaining + return + self._depth.pop(lock_path, None) + descriptor = self._files.pop(lock_path, None) + if descriptor is None: + return + try: + if fcntl is not None: + fcntl.flock(descriptor.fileno(), fcntl.LOCK_UN) + finally: + descriptor.close() + + +_LOCKS = _ReentrantFileLock() + + +@contextmanager +def ingest_lock(persist_path: str | Path) -> Iterator[None]: + """Own the persist path for the duration of a write, or fail fast if someone else does. + + Fails immediately rather than queueing: an ingest can run for minutes, so a caller + blocked behind one is better told to retry than left hanging with no output. + """ + # Resolve first: the in-process depth counter is keyed by path, so two spellings of + # the same directory must not look like two independent locks. + lock_path = Path(persist_path).expanduser().resolve() / LOCK_FILE_NAME + held = _LOCKS.acquire(lock_path) + try: + yield + finally: + if held: + _LOCKS.release(lock_path) diff --git a/localrag/storage/vector_store.py b/localrag/storage/vector_store.py index 826b9b2..d0e87d3 100644 --- a/localrag/storage/vector_store.py +++ b/localrag/storage/vector_store.py @@ -1,6 +1,8 @@ from __future__ import annotations import logging +import shutil +import sqlite3 import threading from dataclasses import dataclass from dataclasses import field as dataclass_field @@ -24,6 +26,7 @@ class VectorStore: client: chromadb.ClientAPI collection: Collection + persist_path: Path | None = None _write_lock: threading.RLock = dataclass_field( default_factory=threading.RLock, init=False, repr=False ) @@ -41,13 +44,25 @@ def create(cls, persist_path: str, collection_name: str) -> VectorStore: persist_path, collection_name, ) - return cls(client=client, collection=collection) + return cls(client=client, collection=collection, persist_path=Path(persist_path)) @classmethod def open(cls, persist_path: str, collection_name: str) -> VectorStore: """Open an existing collection without creating or mutating it.""" client = chromadb.PersistentClient(path=persist_path) - return cls(client=client, collection=client.get_collection(name=collection_name)) + return cls( + client=client, + collection=client.get_collection(name=collection_name), + persist_path=Path(persist_path), + ) + + def for_collection(self, name: str) -> VectorStore: + """Open another collection on this client's persist path.""" + return type(self)( + client=self.client, + collection=self.client.get_collection(name=name), + persist_path=self.persist_path, + ) def add_chunks( self, @@ -333,13 +348,42 @@ def list_collections(self) -> list[str]: def delete_collection(self, name: str) -> None: with self._write_lock: + segment_ids = self._persisted_segment_ids(name) self.client.delete_collection(name) + for segment_id in segment_ids: + segment_path = (self.persist_path / segment_id) if self.persist_path else None + if segment_path is None: + continue + try: + _remove_directory(segment_path) + except OSError: + logger.warning("vector_collection_segment_cleanup_failed path=%s", segment_path) if getattr(self.collection, "name", None) == name: self.collection = self.client.get_or_create_collection( name=name, metadata={"hnsw:space": "cosine"} ) logger.warning("vector_collection_deleted name=%s", name) + def _persisted_segment_ids(self, name: str) -> list[str]: + """Read HNSW segment IDs before Chroma removes their collection metadata.""" + if self.persist_path is None: + return [] + database = self.persist_path / "chroma.sqlite3" + if not database.is_file(): + return [] + try: + collection_id = str(self.client.get_collection(name=name).id) + with sqlite3.connect(database) as connection: + rows = connection.execute( + "SELECT id FROM segments WHERE collection = ? " + "AND type = 'urn:chroma:segment/vector/hnsw-local-persisted'", + (collection_id,), + ).fetchall() + except (OSError, sqlite3.Error): + logger.warning("vector_collection_segment_lookup_failed name=%s", name, exc_info=True) + return [] + return [str(row[0]) for row in rows] + def get_all_chunks(self) -> list[tuple[str, str, dict[str, Any]]]: with self._write_lock: raw = self.collection.get(include=["documents", "metadatas"]) @@ -357,3 +401,9 @@ def get_all_chunks(self) -> list[tuple[str, str, dict[str, Any]]]: @staticmethod def _chunk_id(source: str, chunk_index: int) -> str: return sha1(f"{source}:{chunk_index}".encode(), usedforsecurity=False).hexdigest() + + +def _remove_directory(path: Path) -> None: + """Remove one Chroma segment without touching unrelated persist-path entries.""" + if path.is_dir(): + shutil.rmtree(path) diff --git a/tests/conftest.py b/tests/conftest.py index 883354f..0eee61a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,5 +3,48 @@ from __future__ import annotations import os +import subprocess +import sys +from collections.abc import Callable, Iterator +from pathlib import Path + +import pytest os.environ.setdefault("LOG_LEVEL", "ERROR") + +_HOLD_INGEST_LOCK = """ +import sys +import time + +from localrag.storage.persist_lock import ingest_lock + +with ingest_lock(sys.argv[1]): + print("acquired", flush=True) + time.sleep(float(sys.argv[2])) +""" + +IngestLockHolder = Callable[[Path, float], "subprocess.Popen[str]"] + + +@pytest.fixture +def ingest_lock_holder() -> Iterator[IngestLockHolder]: + """Hold the ingest lock elsewhere; ``flock`` is per-fd, so one process cannot contend.""" + processes: list[subprocess.Popen[str]] = [] + + def spawn(persist_path: Path, hold_seconds: float) -> subprocess.Popen[str]: + process = subprocess.Popen( # noqa: S603 — fixed argv, no shell, test-only helper + [sys.executable, "-c", _HOLD_INGEST_LOCK, str(persist_path), str(hold_seconds)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + processes.append(process) + assert process.stdout is not None + assert process.stdout.readline().strip() == "acquired" + return process + + yield spawn + + for process in processes: + process.kill() + process.wait(timeout=10) diff --git a/tests/integration/test_stack.py b/tests/integration/test_stack.py index dbee3c6..0b87a86 100644 --- a/tests/integration/test_stack.py +++ b/tests/integration/test_stack.py @@ -23,12 +23,19 @@ def test_health_and_readiness(base_url: str) -> None: assert readiness.json() == {"status": "ok"} -def test_metrics_endpoint(base_url: str) -> None: - response = httpx.get(f"{base_url}/metrics", timeout=10.0) +def test_metrics_endpoint(base_url: str, api_key: str) -> None: + response = httpx.get(f"{base_url}/metrics", headers=_headers(api_key), timeout=10.0) assert response.status_code == 200 assert "localrag_query_duration_seconds" in response.text +def test_build_info_requires_auth_and_reports_identity(base_url: str, api_key: str) -> None: + assert httpx.get(f"{base_url}/build-info", timeout=10.0).status_code == 401 + response = httpx.get(f"{base_url}/build-info", headers=_headers(api_key), timeout=10.0) + assert response.status_code == 200 + assert response.json()["build_sha"] + + def test_api_key_missing_returns_401(base_url: str, auth_enabled: bool) -> None: if not auth_enabled: pytest.fail("Auth must be enabled for the integration stack") @@ -80,19 +87,37 @@ def test_query_json_endpoint(base_url: str, auth_enabled: bool, api_key: str) -> pytest.fail("LOCALRAG_TEST_API_KEY is required") response = httpx.post( f"{base_url}/query", - json={"question": "What is LocalRAG?", "model": "qwen2.5:0.5b"}, + json={"question": "What is LocalRAG?", "model": "gemma3:4b"}, headers=_headers(api_key), timeout=60.0, ) assert response.status_code == 200 + collections = httpx.get(f"{base_url}/collections", headers=_headers(api_key), timeout=10.0) + collection = collections.json()["collections"][0] + selected = httpx.post( + f"{base_url}/query", + json={"question": "What is LocalRAG?", "collection": collection}, + headers=_headers(api_key), + timeout=60.0, + ) + assert selected.status_code == 200 + + missing = httpx.post( + f"{base_url}/query", + json={"question": "What is LocalRAG?", "collection": "missing-collection"}, + headers=_headers(api_key), + timeout=10.0, + ) + assert missing.status_code == 404 + def test_query_stream_endpoint(base_url: str, auth_enabled: bool, api_key: str) -> None: if auth_enabled and not api_key: pytest.fail("LOCALRAG_TEST_API_KEY is required") response = httpx.post( f"{base_url}/query/stream", - json={"question": "Give me one sentence about LocalRAG.", "model": "qwen2.5:0.5b"}, + json={"question": "Give me one sentence about LocalRAG.", "model": "gemma3:4b"}, headers=_headers(api_key), timeout=60.0, ) diff --git a/tests/integration/test_vector_store_lifecycle.py b/tests/integration/test_vector_store_lifecycle.py new file mode 100644 index 0000000..2368f91 --- /dev/null +++ b/tests/integration/test_vector_store_lifecycle.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from localrag.storage.vector_store import VectorStore + +pytestmark = pytest.mark.integration + + +def test_real_chroma_delete_reclaims_hnsw_segment(tmp_path: Path) -> None: + persist_path = tmp_path / "chroma" + store = VectorStore.create(str(persist_path), "lifecycle") + count = 1001 + store.add_chunks( + source="fixture", + chunks=["fixture text"] * count, + embeddings=[[1.0, 0.0, 0.0]] * count, + metadatas=[{"chunk_id": str(index), "source": "fixture"} for index in range(count)], + ) + segment_directories = [path for path in persist_path.iterdir() if path.is_dir()] + assert segment_directories + + store.delete_collection("lifecycle") + + assert not any(path.exists() for path in segment_directories) diff --git a/tests/test_api.py b/tests/test_api.py index fdd97c8..46ab915 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, field +from http import HTTPStatus from pathlib import Path from typing import Any from uuid import uuid4 @@ -16,6 +17,7 @@ from localrag.api.main import app from localrag.ingestion.service import IngestionResult from localrag.settings import Settings, get_settings +from localrag.storage.persist_lock import ConcurrentIngestError _STUB_CONTEXTS = [{"source": "doc.md", "chunk_index": 1, "text": "chunk"}] @@ -85,6 +87,23 @@ def test_query_json_returns_answer() -> None: app.dependency_overrides.clear() +def test_query_json_selects_request_collection() -> None: + selected: list[str] = [] + + @dataclass + class CollectionEngine(StubEngine): + def for_collection(self, collection: str) -> CollectionEngine: + selected.append(collection) + return self + + app.dependency_overrides[get_engine] = lambda: CollectionEngine() + response = TestClient(app).post("/query", json={"question": "Hi", "collection": "experiments"}) + + assert response.status_code == 200 + assert selected == ["experiments"] + app.dependency_overrides.clear() + + def test_benchmark_contexts_return_text_and_stable_id() -> None: class BenchmarkRetriever(StubRetriever): def retrieve(self, **_kwargs: object) -> list[dict[str, Any]]: @@ -138,6 +157,18 @@ def test_metrics_endpoint_returns_prometheus_text() -> None: assert "python_info" in response.text or "HELP" in response.text +def test_build_info_is_protected(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LOCALRAG_BUILD_SHA", "test-sha") + app.dependency_overrides[get_settings] = lambda: Settings(api_key="secret") + client = TestClient(app) + + assert client.get("/build-info").status_code == 401 + response = client.get("/build-info", headers={"X-API-Key": "secret"}) + + assert response.json() == {"build_sha": "test-sha"} + app.dependency_overrides.clear() + + @pytest.mark.parametrize( ("headers", "expected_status"), [ @@ -340,3 +371,23 @@ def ingest_file(self, path: Path, embed_model: str | None = None) -> IngestionRe assert old.exists() is False assert len(list(tmp_path.iterdir())) == 1 app.dependency_overrides.clear() + + +def test_ingest_returns_409_when_another_process_holds_the_persist_lock(tmp_path: Path) -> None: + doc = tmp_path / "notes.txt" + doc.write_text("hello", encoding="utf-8") + + @dataclass + class ContendedIngestionService: + def ingest_file(self, path: Path, embed_model: str | None = None) -> IngestionResult: + _ = embed_model + raise ConcurrentIngestError(str(path)) + + app.dependency_overrides[get_ingestion_service] = lambda: ContendedIngestionService() + client = TestClient(app, raise_server_exceptions=False) + + response = client.post("/ingest", json={"path": str(doc)}) + + assert response.status_code == HTTPStatus.CONFLICT + assert str(doc) in response.json()["detail"] + app.dependency_overrides.clear() diff --git a/tests/test_ingestion_service.py b/tests/test_ingestion_service.py index 35ff1c9..7bd3ea6 100644 --- a/tests/test_ingestion_service.py +++ b/tests/test_ingestion_service.py @@ -1,7 +1,8 @@ from __future__ import annotations import hashlib -from collections.abc import Sequence +import subprocess +from collections.abc import Callable, Sequence from dataclasses import dataclass, field from pathlib import Path @@ -11,6 +12,9 @@ from localrag.ingestion import service as service_module from localrag.ingestion.service import FailedSource, IngestionResult, IngestionService from localrag.settings import Settings +from localrag.storage.persist_lock import ConcurrentIngestError + +IngestLockHolder = Callable[[Path, float], "subprocess.Popen[str]"] @dataclass @@ -577,3 +581,66 @@ def embed_texts( assert ("a.md", None) in events assert any(name == "b.md" and error is not None for name, error in events) + + +@pytest.mark.parametrize( + "call_name", + ["ingest_file", "ingest_directory", "ingest_paths", "rebuild_collection"], +) +def test_ingest_entry_points_reject_a_concurrent_process( + tmp_path: Path, call_name: str, ingest_lock_holder: IngestLockHolder +) -> None: + document = tmp_path / "a.md" + document.write_text("hello world", encoding="utf-8") + persist_path = tmp_path / "chroma" + settings = Settings( + ingest_roots=[str(tmp_path)], + chroma_persist_path=str(persist_path), + chunk_chars=100, + chunk_overlap_chars=0, + ) + service = IngestionService( + settings=settings, + embedder=StubEmbedder(seen_texts_batches=[]), # type: ignore[arg-type] + vector_store=StubVectorStore(deleted_sources=[], added=[]), # type: ignore[arg-type] + ) + calls: dict[str, Callable[[], object]] = { + "ingest_file": lambda: service.ingest_file(document), + "ingest_directory": lambda: service.ingest_directory(tmp_path), + "ingest_paths": lambda: service.ingest_paths([document]), + "rebuild_collection": service.rebuild_collection, + } + + ingest_lock_holder(persist_path, 30.0) + + with pytest.raises(ConcurrentIngestError): + calls[call_name]() + + +def test_ingest_succeeds_once_the_competing_process_has_finished( + tmp_path: Path, ingest_lock_holder: IngestLockHolder +) -> None: + document = tmp_path / "a.md" + document.write_text("hello world", encoding="utf-8") + persist_path = tmp_path / "chroma" + settings = Settings( + ingest_roots=[str(tmp_path)], + chroma_persist_path=str(persist_path), + chunk_chars=100, + chunk_overlap_chars=0, + ) + vector_store = StubVectorStore(deleted_sources=[], added=[]) + service = IngestionService( + settings=settings, + embedder=StubEmbedder(seen_texts_batches=[]), # type: ignore[arg-type] + vector_store=vector_store, # type: ignore[arg-type] + ) + + ingest_lock_holder(persist_path, 0.0).wait(timeout=10) + + first = service.ingest_file(document) + second = service.ingest_file(document) + + assert first.files_processed == 1 + assert second.files_processed == 1 + assert len(vector_store.added) == 2 diff --git a/tests/test_persist_lock.py b/tests/test_persist_lock.py new file mode 100644 index 0000000..327da48 --- /dev/null +++ b/tests/test_persist_lock.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import subprocess +from collections.abc import Callable +from pathlib import Path + +import pytest + +from localrag.storage import persist_lock +from localrag.storage.persist_lock import ConcurrentIngestError, ingest_lock + +IngestLockHolder = Callable[[Path, float], "subprocess.Popen[str]"] + + +def test_second_holder_is_rejected_while_the_first_still_holds_the_lock( + tmp_path: Path, ingest_lock_holder: IngestLockHolder +) -> None: + persist_path = tmp_path / "chroma" + ingest_lock_holder(persist_path, 30.0) + + with pytest.raises(ConcurrentIngestError) as excinfo, ingest_lock(persist_path): + pytest.fail("the lock must not be granted twice") + + assert str(persist_path) in str(excinfo.value) + + +def test_lock_is_reacquirable_after_the_holder_exits( + tmp_path: Path, ingest_lock_holder: IngestLockHolder +) -> None: + persist_path = tmp_path / "chroma" + ingest_lock_holder(persist_path, 0.0).wait(timeout=10) + + with ingest_lock(persist_path): + pass + + with ingest_lock(persist_path): + pass + + +def test_nested_acquisition_keeps_the_lock_until_the_outermost_scope_exits( + tmp_path: Path, ingest_lock_holder: IngestLockHolder +) -> None: + """Rebuild delegates to the ingest path, so the boundary must tolerate nesting.""" + persist_path = tmp_path / "chroma" + + with ingest_lock(persist_path), ingest_lock(persist_path): + pass + + ingest_lock_holder(persist_path, 30.0) + with pytest.raises(ConcurrentIngestError), ingest_lock(persist_path): + pytest.fail("an inner release must not hand the lock to another process") + + +def test_missing_fcntl_degrades_to_a_no_op(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + persist_path = tmp_path / "chroma" + monkeypatch.setattr(persist_lock, "fcntl", None) + + with ingest_lock(persist_path), ingest_lock(persist_path): + pass diff --git a/tests/test_retriever.py b/tests/test_retriever.py index d0f0523..bc74241 100644 --- a/tests/test_retriever.py +++ b/tests/test_retriever.py @@ -57,6 +57,30 @@ def test_retriever_returns_contexts() -> None: ] +def test_retriever_for_collection_creates_request_scoped_store() -> None: + @dataclass + class CollectionStore(StubStore): + selected: list[str] + + def for_collection(self, name: str) -> CollectionStore: + self.selected.append(name) + return self + + settings = Settings(chroma_collection_name="default") + store = CollectionStore(selected=[]) + retriever = Retriever( + settings=settings, + embedder=StubEmbedder(), # type: ignore[arg-type] + vector_store=store, # type: ignore[arg-type] + ) + + selected = retriever.for_collection("experiments") + + assert store.selected == ["experiments"] + assert selected.settings.chroma_collection_name == "experiments" + assert selected.embedder is retriever.embedder + + @respx.mock def test_retriever_raises_retrieval_failure_when_ollama_embed_fails() -> None: respx.post("http://ollama:11434/api/embed").mock(return_value=httpx.Response(503)) diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index c375dc6..94bd4c2 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -1,8 +1,10 @@ from __future__ import annotations +import sqlite3 from dataclasses import dataclass, field from hashlib import sha1 from pathlib import Path +from types import SimpleNamespace import pytest @@ -385,6 +387,47 @@ def test_vector_store_get_chunks_by_headings_filters_and_groups_in_one_lookup() assert sections == {("guide.md", "Setup"): [(1, "team-a first"), (2, "team-a second")]} +def test_delete_collection_removes_only_its_persisted_hnsw_segment(tmp_path: Path) -> None: + persist_path = tmp_path / "chroma" + persist_path.mkdir() + database = persist_path / "chroma.sqlite3" + collection_id = "collection-id" + segment_id = "segment-id" + unrelated_segment_id = "unrelated-segment-id" + connection = sqlite3.connect(database) + connection.execute("CREATE TABLE segments (id TEXT, type TEXT, collection TEXT)") + connection.execute( + "INSERT INTO segments VALUES (?, ?, ?)", + (segment_id, "urn:chroma:segment/vector/hnsw-local-persisted", collection_id), + ) + connection.commit() + connection.close() + (persist_path / segment_id).mkdir() + (persist_path / unrelated_segment_id).mkdir() + + collection = SimpleNamespace(name="target", id=collection_id) + + class Client: + def get_collection(self, name: str) -> object: + assert name == "target" + return collection + + def delete_collection(self, name: str) -> None: + assert name == "target" + + def get_or_create_collection(self, name: str, metadata: dict[str, str]) -> object: + assert name == "target" + _ = metadata + return collection + + store = VectorStore(client=Client(), collection=collection, persist_path=persist_path) # type: ignore[arg-type] + + store.delete_collection("target") + + assert not (persist_path / segment_id).exists() + assert (persist_path / unrelated_segment_id).exists() + + def test_vector_store_create_initializes_persistent_client( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: