From af37c6e4d187ac2a3677692d6fb4ca9c5b299774 Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 19:40:00 +0200 Subject: [PATCH 1/8] refactor(storage): a VectorStore interface with the clearance guard in the base class Retrieval, ingest and workspace management no longer talk to Chroma directly: they go through VectorStore (list, ensure, exists, count, metadata, read_vectors, publish, delete, query). Chroma is the default and behaves exactly as before (same collection names, metadata, atomic swap by renaming, the unfiltered fallback for old collections), which the unchanged suite confirms. VectorStore.query() applies guard_hits() after the adapter's server-side filter: a chunk without a valid clearance, with one the caller may not read, or from another workspace is dropped. A broken filter in an adapter cannot leak, and a new adapter cannot forget it; a test with a deliberately leaky adapter pins that. Isolation: the new scaling settings are pinned off in tests and in the eval redirect, and an explicitly empty PGVECTOR_URI means off (the guard caught the fallback that built one from POSTGRES_PASSWORD). --- rag_lab/config.py | 46 ++++- rag_lab/ingestion/ingest.py | 86 +------- rag_lab/ops/eval.py | 8 + rag_lab/pipeline/retrieval.py | 53 +---- rag_lab/storage/vector_store.py | 282 ++++++++++++++++++++++++++ rag_lab/storage/workspace_manager.py | 40 ++-- tests/conftest.py | 12 ++ tests/test_answers.py | 2 +- tests/test_ingest_reuse.py | 4 +- tests/test_rbac.py | 2 +- tests/test_vector_stores.py | 291 +++++++++++++++++++++++++++ 11 files changed, 676 insertions(+), 150 deletions(-) create mode 100644 rag_lab/storage/vector_store.py create mode 100644 tests/test_vector_stores.py diff --git a/rag_lab/config.py b/rag_lab/config.py index fc50cc7..73909bc 100644 --- a/rag_lab/config.py +++ b/rag_lab/config.py @@ -11,7 +11,7 @@ from __future__ import annotations import os -from dataclasses import dataclass, fields +from dataclasses import dataclass, field, fields from dotenv import load_dotenv @@ -25,6 +25,18 @@ DEFAULT_DB_URI = f"sqlite:///{DEFAULT_SQLITE_PATH}" +def pgvector_uri_from_env() -> str: + """The pgvector compose profile's server (loopback, port 5433), with the same credentials as the checkpoint + database. Empty when POSTGRES_PASSWORD is not set: there is no default password.""" + from urllib.parse import quote + + user, password = os.getenv("POSTGRES_USER", "rag_user"), os.getenv("POSTGRES_PASSWORD", "") + if not password: + return "" + port = os.getenv("PGVECTOR_PORT", "5433") + return f"postgresql://{quote(user, safe='')}:{quote(password, safe='')}@localhost:{port}/rag_vectors" + + def postgres_uri_from_env() -> str: """The URI of the docker-compose PostgreSQL server, built from POSTGRES_USER / POSTGRES_PASSWORD (.env). @@ -141,6 +153,26 @@ class Settings: # Whether `make recover` (the operator-run break-glass for a locked-out account) may run. On by default: this # is a local-first tool and its operator owns the machine. Turn it off on a shared host. local_recovery: bool + # --- scaling options, all off by default (docs/scaling.md) --------------------------------------------------- + # Where chunk vectors live: "chroma" (embedded, default), "qdrant" or "pgvector". See rag_lab/storage/vector_store.py. + vector_store: str + qdrant_url: str + # Secrets and connection strings with passwords stay out of repr(), so printing settings cannot leak them. + qdrant_api_key: str = field(repr=False) + # Prepended to every Qdrant collection name: lets installs (and tests) share one server without touching each other. + qdrant_prefix: str + pgvector_uri: str = field(repr=False) + # The schema holding rag_chunks / rag_workspaces in the pgvector database. + pgvector_schema: str + # Where the semantic cache lives: "sqlite" (default, a local file) or "postgres" (shared by several UI processes). + cache_backend: str + cache_db_uri: str = field(repr=False) + # OpenAI-compatible chat server (vLLM, LM Studio, llama.cpp), used with LLM_PROVIDER=openai. + local_llm_api_key: str = field(repr=False) + local_llm_model: str + # JSON merged into the request as chat_template_kwargs. The default turns Qwen3's thinking off, as the Ollama path + # does; set it empty for a server that rejects unknown request fields. + local_llm_template_kwargs: str # The native accelerator (rag_accel): "preferred" uses it and falls back, loudly, to pure Python when it is # missing; "required" refuses to start without it. See rag_lab/accel.py. rag_accel: str @@ -225,6 +257,18 @@ def from_env(cls) -> Settings: key_envelope=_env("KEY_ENVELOPE", "").strip(), local_recovery=_flag("LOCAL_RECOVERY", "on"), rag_accel=_env("RAG_ACCEL", "preferred").strip().lower(), + vector_store=_env("VECTOR_STORE", "chroma").strip().lower(), + qdrant_url=_env("QDRANT_URL", "http://localhost:6333"), + qdrant_api_key=_env("QDRANT_API_KEY", ""), + qdrant_prefix=_env("QDRANT_PREFIX", ""), + # Unset: derive it from POSTGRES_* (the compose profile). Set but empty: off, however .env is filled. + pgvector_uri=os.environ["PGVECTOR_URI"] if "PGVECTOR_URI" in os.environ else pgvector_uri_from_env(), + pgvector_schema=_env("PGVECTOR_SCHEMA", "public"), + cache_backend=_env("CACHE_BACKEND", "sqlite").strip().lower(), + cache_db_uri=_env("CACHE_DB_URI", ""), + local_llm_api_key=_env("LOCAL_LLM_API_KEY", "local"), + local_llm_model=_env("LOCAL_LLM_MODEL", ""), + local_llm_template_kwargs=_env("LOCAL_LLM_TEMPLATE_KWARGS", '{"enable_thinking": false}'), backup_dir=_env("BACKUP_DIR", str(ROOT / "backups")), retention_days=int(_env("RETENTION_DAYS", "183")), max_source_bytes=int(_env("MAX_SOURCE_BYTES", str(25 * 1024 * 1024))), diff --git a/rag_lab/ingestion/ingest.py b/rag_lab/ingestion/ingest.py index c5380d8..fea1cc8 100644 --- a/rag_lab/ingestion/ingest.py +++ b/rag_lab/ingestion/ingest.py @@ -19,7 +19,6 @@ from dataclasses import dataclass, field from pathlib import Path -import chromadb import numpy as np import yaml from langchain_text_splitters import RecursiveCharacterTextSplitter @@ -30,7 +29,6 @@ from rag_lab.config import settings from rag_lab.security.context_sanitizer import scan_chunk from rag_lab.retrieval.embedding_provider import ( - FINGERPRINT_KEY, EmbeddingUnavailableError, embedding_fingerprint, get_embedding_model, @@ -38,19 +36,11 @@ ) from rag_lab.ingestion.parsers import SUPPORTED_EXTENSIONS, extract_text_from_file, unsupported_reason from rag_lab.retrieval.sparse_search import delete_sparse_index, has_sparse_index, index_sparse_chunks +from rag_lab.storage.vector_store import get_vector_store from rag_lab.storage.workspace_manager import ensure_base_dirs, sanitize_workspace_id logger = logging.getLogger("ingest") -# Chroma rejects a single add() above a few thousand rows. -_ADD_BATCH = 256 -# Cosine distance keeps similarity meaningful whichever embedding model produced the vectors. -_COLLECTION_METADATA = {"hnsw:space": "cosine"} - - -def _collection_metadata(fingerprint: str) -> dict[str, str]: - """Metadata for a new collection: distance space plus the model that built its vectors.""" - return {**_COLLECTION_METADATA, FINGERPRINT_KEY: fingerprint} _FRONT_MATTER = re.compile(r"\A---[ \t]*\r?\n(.*?)\r?\n---[ \t]*(?:\r?\n|\Z)", re.DOTALL) _HEADER_TAG = re.compile(r"\A(?:#\s*clearance:\s*|\[clearance\]:\s*)([a-zA-Z0-9_-]+)[ \t]*(?:\r?\n|\Z)", re.IGNORECASE) @@ -125,20 +115,14 @@ def _reusable_vectors(ws_id: str, fingerprint: str, manifest: dict | None, force if manifest.get("embedding") != fingerprint or manifest.get("passage_prefix") != settings.embedding_passage_prefix: return {} try: - collection = chromadb.PersistentClient(path=settings.chroma_path).get_collection(f"ws_{ws_id}") - if stored_fingerprint(getattr(collection, "metadata", None), manifest) != fingerprint: + store = get_vector_store() + if stored_fingerprint(store.collection_metadata(ws_id), manifest) != fingerprint: return {} found: dict[str, list[float]] = {} - total = collection.count() - for offset in range(0, total, 1000): - page = collection.get(include=["documents", "embeddings"], limit=1000, offset=offset) - texts, vectors = page.get("documents") or [], page.get("embeddings") - if vectors is None or len(texts) != len(vectors): - return {} - for text, vector in zip(texts, vectors): - array = np.asarray(vector, dtype=np.float64) - if text and array.ndim == 1 and array.size and bool(np.isfinite(array).all()): - found.setdefault(text, array.tolist()) + for text, vector in store.read_vectors(ws_id): + array = np.asarray(vector, dtype=np.float64) + if text and array.ndim == 1 and array.size and bool(np.isfinite(array).all()): + found.setdefault(text, array.tolist()) return found except Exception: logger.debug("Could not read the previous vectors of %s; embedding everything", ws_id, exc_info=True) @@ -172,54 +156,9 @@ def _save_manifest(ws_id: str, manifest_data: dict) -> None: os.replace(temporary, path) -def _add_in_batches(collection, ids, documents, metadatas, embeddings) -> None: - for start in range(0, len(ids), _ADD_BATCH): - end = start + _ADD_BATCH - collection.add( - ids=ids[start:end], - documents=documents[start:end], - metadatas=metadatas[start:end], - embeddings=embeddings[start:end], - ) - - -def _delete_quiet(client, name: str) -> None: - try: - client.delete_collection(name=name) - except Exception as exc: - logger.info("Collection %s was not deleted (%s)", name, exc) - - -def _publish_collection(client, name: str, ids, documents, metadatas, embeddings, fingerprint: str) -> None: - """Build the new collection beside the live one, then swap them by renaming. - - Readers see the old index until the swap, and the gap while it happens is two renames long. - If the swap fails the old collection is put back, so a failed publish never loses the live index. - """ - build_name, old_name = f"{name}__build", f"{name}__old" - _delete_quiet(client, build_name) - _delete_quiet(client, old_name) - building = client.create_collection(name=build_name, metadata=_collection_metadata(fingerprint)) - _add_in_batches(building, ids, documents, metadatas, embeddings) - - live_exists = name in [getattr(c, "name", str(c)) for c in client.list_collections()] - if live_exists: - client.get_collection(name).modify(name=old_name) - try: - building.modify(name=name) - except Exception: - logger.exception("Could not publish collection %s; restoring the previous index", name) - if live_exists: - client.get_collection(old_name).modify(name=name) - raise - _delete_quiet(client, old_name) - - def _clear_index(ws_id: str) -> None: """Remove every index artifact of a workspace: vectors, keywords, manifest, and cached answers.""" - client = chromadb.PersistentClient(path=settings.chroma_path) - for name in (f"ws_{ws_id}", f"ws_{ws_id}__build", f"ws_{ws_id}__old"): - _delete_quiet(client, name) + get_vector_store().delete(ws_id) delete_sparse_index(ws_id, base_dir=settings.bm25_dir) manifest = _get_manifest_path(ws_id) if manifest.exists(): @@ -438,10 +377,8 @@ def ingest_workspace( and manifest.get("chunker") == CHUNKER_VERSION ): try: - client = chromadb.PersistentClient(path=settings.chroma_path) - existing_cols = [getattr(c, "name", str(c)) for c in client.list_collections()] fts_exists = os.path.exists(os.path.join(settings.bm25_dir, f"{ws_id}.db")) - if f"ws_{ws_id}" in existing_cols and fts_exists: + if get_vector_store().exists(ws_id) and fts_exists: result.documents = len(raw_docs) result.chunks = manifest.get("total_chunks", 0) result.indexed_files = list(file_hashes.keys()) @@ -504,12 +441,9 @@ def ingest_workspace( result.reused = sum(1 for text in chunks if text in stored) result.embedded = len(to_embed) - chroma_path = settings.chroma_path bm25_dir = settings.bm25_dir - collection_name = f"ws_{ws_id}" try: - client = chromadb.PersistentClient(path=chroma_path) - _publish_collection(client, collection_name, chunk_ids, chunks, metadatas, vectors, fingerprint) + get_vector_store().publish(ws_id, chunk_ids, chunks, metadatas, vectors, fingerprint) except Exception as exc: logger.exception("Chroma write failed for workspace %s", ws_id) result.error = f"Vector index write failed: {exc}" diff --git a/rag_lab/ops/eval.py b/rag_lab/ops/eval.py index c82a8ce..0572b4a 100644 --- a/rag_lab/ops/eval.py +++ b/rag_lab/ops/eval.py @@ -62,6 +62,14 @@ def redirect_stores(out_dir: Path, llm_mode: bool) -> None: os.environ["KEY_STORE"] = "file" os.environ["KEY_ENVELOPE"] = "" os.environ["AUDIT_ANCHOR_COPY_DIR"] = "" + # The evaluation measures the embedded defaults; a .env that enables a scaling option must not redirect it. + # To evaluate an adapter on purpose, say so in the shell: `VECTOR_STORE=qdrant make eval` (docs/scaling.md). Even + # then it gets its own namespace, never the real collections or schema. + os.environ.setdefault("VECTOR_STORE", "chroma") + os.environ.setdefault("QDRANT_PREFIX", "eval_") + os.environ.setdefault("PGVECTOR_SCHEMA", "eval") + os.environ["CACHE_BACKEND"] = "sqlite" + os.environ["CACHE_DB_URI"] = "" os.environ["DB_URI"] = f"sqlite:///{out_dir / 'unused_checkpoints.db'}" if llm_mode: os.environ["LLM_PROVIDER"] = "ollama" diff --git a/rag_lab/pipeline/retrieval.py b/rag_lab/pipeline/retrieval.py index 8ef8e4b..2dc918b 100644 --- a/rag_lab/pipeline/retrieval.py +++ b/rag_lab/pipeline/retrieval.py @@ -10,8 +10,6 @@ import json import time -from chromadb.errors import NotFoundError -import chromadb from rag_lab.config import settings from rag_lab.pipeline import llm_provider @@ -20,6 +18,7 @@ from rag_lab.retrieval.embedding_provider import embedding_fingerprint from rag_lab.retrieval.embedding_provider import get_embedding_model from rag_lab.retrieval.embedding_provider import stored_fingerprint +from rag_lab.storage.vector_store import WorkspaceNotIndexed, get_vector_store from rag_lab.retrieval.reranker import rerank_with_status, score_pairs from rag_lab.retrieval.sparse_search import has_sparse_index from rag_lab.retrieval.sparse_search import search_sparse_chunks @@ -34,17 +33,6 @@ from rag_lab.pipeline.text import _content_terms, _coverage, _keywords, _stem_variants -def _similarity(distance: float, space: str) -> float: - """Convert a Chroma distance to a cosine-like similarity for the collection's distance space. - - Collections built by this version use cosine distance. Older ones use squared L2 on unit vectors, - where cosine = 1 - d / 2. - """ - if space == "cosine": - return 1.0 - distance - return 1.0 - distance / 2.0 - - class _EmbeddingMismatch(RuntimeError): """Raised inside dense retrieval to skip it with a labelled note.""" @@ -146,41 +134,22 @@ def hybrid_retrieval_node(state: StateInput) -> dict[str, Any]: user_clr = get_user_clearance(principal) if principal else "public" authorized_clrs = allowed_clearances(user_clr) - # 1. Dense retrieval (ChromaDB) with clearance filter + fail-closed post-query verification + # 1. Dense retrieval through the configured vector store (VECTOR_STORE): clearance filter on the server, and the + # fail-closed guard of VectorStore.query() on every hit, whichever store it is. try: - chroma_client = chromadb.PersistentClient(path=settings.chroma_path) - collection = chroma_client.get_collection(name=f"ws_{ws_id}") - raw_meta = getattr(collection, "metadata", None) - space = str(raw_meta.get("hnsw:space", "l2")) if isinstance(raw_meta, dict) else "l2" - mismatch = _embedding_mismatch(ws_id, raw_meta) + store = get_vector_store() + mismatch = _embedding_mismatch(ws_id, store.collection_metadata(ws_id)) if mismatch: raise _EmbeddingMismatch(mismatch) q_vec = get_embedding_model().embed_query(query) - try: - results = collection.query( - query_embeddings=[q_vec], # type: ignore[arg-type] # chroma stubs want numpy; lists are accepted - n_results=5, - where={"clearance": {"$in": authorized_clrs}}, # type: ignore[dict-item] - ) - except Exception: - results = collection.query(query_embeddings=[q_vec], n_results=20) # type: ignore[arg-type] - docs = (results.get("documents") or [[]])[0] - metas = (results.get("metadatas") or [[]])[0] or [{}] * len(docs) - ids = (results.get("ids") or [[]])[0] - distances = (results.get("distances") or [[]])[0] or [2.0] * len(docs) - for doc, meta, cid, distance in zip(docs, metas, ids, distances): - meta_dict = meta or {} - # A chunk with no clearance tag is treated as unreadable, never as public. - doc_clr = str(meta_dict.get("clearance", "")).strip().lower() - if not doc_clr or doc_clr not in authorized_clrs: - continue + for hit in store.query(ws_id, q_vec, 5, authorized_clrs): dense_candidates.append({ - "id": cid, - "content": doc, - "citation": f"[{cid}] {meta_dict.get('source', 'document')}", - "similarity": _similarity(float(distance), space), + "id": hit.id, + "content": hit.document, + "citation": f"[{hit.id}] {hit.metadata.get('source', 'document')}", + "similarity": hit.similarity, }) - except NotFoundError: + except WorkspaceNotIndexed: degraded.append("not_indexed: this workspace has no vector index yet") except _EmbeddingMismatch as exc: degraded.append(str(exc)) diff --git a/rag_lab/storage/vector_store.py b/rag_lab/storage/vector_store.py new file mode 100644 index 0000000..f784733 --- /dev/null +++ b/rag_lab/storage/vector_store.py @@ -0,0 +1,282 @@ +"""Where chunk vectors live: embedded Chroma by default, Qdrant or pgvector when a deployment outgrows one host. + +Every store answers the same few questions per workspace (is it indexed, with which model, give me the nearest +chunks the caller may read, replace the whole index atomically), so the pipeline never learns which one it talks to. + +Two rules hold for every adapter, and are enforced here rather than left to each of them: + +* **The clearance guard is in the base class.** `VectorStore.query()` calls the adapter's `_search()` (which should + filter on the server) and then applies `guard_hits()`: a chunk with no valid clearance tag, a clearance the caller + may not read, or a workspace other than the one asked about is dropped. A buggy or compromised filter in an adapter + therefore cannot leak a chunk, and a new adapter cannot forget the guard. +* **Similarity means cosine similarity on the same scale as Chroma's `1 - distance`.** The relevance gates are + calibrated on that scale (make calibrate-models); a store with a different score scale would break them without any + error, as the Ollama embedder did. tests/test_vector_stores.py asserts the values, not just the order. + +Chroma is the default and its behaviour is unchanged: same collection names (`ws_`), the same metadata, the same +atomic swap by renaming a `__build` collection. Qdrant and pgvector are optional (`VECTOR_STORE=qdrant|pgvector`, +docs/scaling.md). +""" +from __future__ import annotations + +import logging +from abc import ABC, abstractmethod +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import Any + +from rag_lab.config import settings +from rag_lab.retrieval.embedding_provider import FINGERPRINT_KEY + +logger = logging.getLogger("vector_store") + +STORES = ("chroma", "qdrant", "pgvector") +_TEMP_SUFFIXES = ("__build", "__old") + + +class VectorStoreError(RuntimeError): + """The configured store is unknown, unreachable or not installed. The message says how to fix it.""" + + +class WorkspaceNotIndexed(LookupError): + """The store has no index for this workspace.""" + + +@dataclass(frozen=True) +class Hit: + id: str + document: str + metadata: dict[str, Any] + similarity: float + + +def guard_hits(hits: Sequence[Hit], workspace: str, allowed_clearances: Sequence[str]) -> list[Hit]: + """Fail-closed check applied to every hit of every adapter after the server-side filter. + + A chunk with no clearance tag is treated as unreadable, never as public. + """ + allowed = {c.strip().lower() for c in allowed_clearances} + kept: list[Hit] = [] + for hit in hits: + meta = hit.metadata or {} + clearance = str(meta.get("clearance", "")).strip().lower() + if not clearance or clearance not in allowed: + continue + if meta.get("workspace", workspace) != workspace: + continue + kept.append(hit) + return kept + + +class VectorStore(ABC): + name = "" + + # --- what an adapter implements ----------------------------------------------------------------- + @abstractmethod + def workspaces(self) -> set[str]: + """Workspaces that have an index (or were created empty), without the temporary copies of a re-index.""" + + @abstractmethod + def ensure(self, workspace: str) -> None: + """Create an empty index for a new workspace if there is none.""" + + @abstractmethod + def exists(self, workspace: str) -> bool: ... + + @abstractmethod + def count(self, workspace: str) -> int: + """Chunks in the workspace's index. Raises WorkspaceNotIndexed when it has none.""" + + @abstractmethod + def collection_metadata(self, workspace: str) -> dict[str, Any] | None: + """What the index records about itself: at least FINGERPRINT_KEY when known. Raises WorkspaceNotIndexed.""" + + @abstractmethod + def read_vectors(self, workspace: str) -> Iterator[tuple[str, list[float]]]: + """(chunk text, vector) for every chunk, so an unchanged chunk need not be embedded again.""" + + @abstractmethod + def publish( + self, + workspace: str, + ids: list[str], + documents: list[str], + metadatas: list[dict[str, Any]], + vectors: list[list[float]], + fingerprint: str, + ) -> None: + """Replace the workspace's index with these chunks atomically: readers see the old index or the new one, and a + failure leaves the old one serving.""" + + @abstractmethod + def delete(self, workspace: str) -> None: ... + + @abstractmethod + def _search(self, workspace: str, vector: list[float], n: int, allowed_clearances: Sequence[str]) -> list[Hit]: + """Nearest chunks, best first, filtered on the server to the allowed clearances where the store can.""" + + # --- what every store shares --------------------------------------------------------------------- + def query(self, workspace: str, vector: list[float], n: int, allowed_clearances: Sequence[str]) -> list[Hit]: + """Nearest chunks the caller may read. Raises WorkspaceNotIndexed when the workspace has no index.""" + return guard_hits(self._search(workspace, vector, n, allowed_clearances), workspace, allowed_clearances) + + +# --- Chroma (the default) -------------------------------------------------------------------------------------- + +def _similarity(distance: float, space: str) -> float: + """Convert a Chroma distance to a cosine-like similarity for the collection's distance space. + + Collections built by this version use cosine distance. Older ones use squared L2 on unit vectors, + where cosine = 1 - d / 2. + """ + if space == "cosine": + return 1.0 - distance + return 1.0 - distance / 2.0 + + +class ChromaStore(VectorStore): + name = "chroma" + BATCH = 256 + # Cosine distance keeps similarity meaningful whichever embedding model produced the vectors. + METADATA = {"hnsw:space": "cosine"} + + def _client(self): + import chromadb + + return chromadb.PersistentClient(path=settings.chroma_path) + + @staticmethod + def _name(workspace: str) -> str: + return f"ws_{workspace}" + + def _collection(self, workspace: str): + from chromadb.errors import NotFoundError + + try: + return self._client().get_collection(name=self._name(workspace)) + except NotFoundError as exc: + raise WorkspaceNotIndexed(workspace) from exc + + def workspaces(self) -> set[str]: + found = set() + for col in self._client().list_collections(): + name = getattr(col, "name", str(col)) + if name.startswith("ws_") and not name.endswith(_TEMP_SUFFIXES): + found.add(name[3:]) + return found + + def ensure(self, workspace: str) -> None: + self._client().get_or_create_collection(name=self._name(workspace), metadata=dict(self.METADATA)) + + def exists(self, workspace: str) -> bool: + return self._name(workspace) in [getattr(c, "name", str(c)) for c in self._client().list_collections()] + + def count(self, workspace: str) -> int: + return int(self._collection(workspace).count()) + + def collection_metadata(self, workspace: str) -> dict[str, Any] | None: + meta = getattr(self._collection(workspace), "metadata", None) + return dict(meta) if isinstance(meta, dict) else None + + def read_vectors(self, workspace: str) -> Iterator[tuple[str, list[float]]]: + collection = self._collection(workspace) + total = collection.count() + for offset in range(0, total, 1000): + page = collection.get(include=["documents", "embeddings"], limit=1000, offset=offset) + texts, vectors = page.get("documents") or [], page.get("embeddings") + if vectors is None or len(texts) != len(vectors): + raise VectorStoreError(f"the index of '{workspace}' returned documents without matching vectors") + for text, vector in zip(texts, vectors): + yield text, [float(x) for x in vector] + + def publish(self, workspace, ids, documents, metadatas, vectors, fingerprint) -> None: + """Build the new collection beside the live one, then swap them by renaming. + + Readers see the old index until the swap, and the gap while it happens is two renames long. If the swap fails + the old collection is put back, so a failed publish never loses the live index. + """ + client = self._client() + name = self._name(workspace) + build_name, old_name = f"{name}__build", f"{name}__old" + self._delete_quiet(client, build_name) + self._delete_quiet(client, old_name) + building = client.create_collection(name=build_name, metadata={**self.METADATA, FINGERPRINT_KEY: fingerprint}) + for start in range(0, len(ids), self.BATCH): + end = start + self.BATCH + building.add(ids=ids[start:end], documents=documents[start:end], metadatas=metadatas[start:end], + embeddings=vectors[start:end]) + + live_exists = name in [getattr(c, "name", str(c)) for c in client.list_collections()] + if live_exists: + client.get_collection(name).modify(name=old_name) + try: + building.modify(name=name) + except Exception: + logger.exception("Could not publish collection %s; restoring the previous index", name) + if live_exists: + client.get_collection(old_name).modify(name=name) + raise + self._delete_quiet(client, old_name) + + def delete(self, workspace: str) -> None: + client = self._client() + name = self._name(workspace) + for target in (name, f"{name}__build", f"{name}__old"): + self._delete_quiet(client, target) + + @staticmethod + def _delete_quiet(client, name: str) -> None: + try: + client.delete_collection(name=name) + except Exception as exc: + logger.debug("Collection %s was not deleted (%s)", name, exc) + + def _search(self, workspace, vector, n, allowed_clearances) -> list[Hit]: + collection = self._collection(workspace) + raw_meta = getattr(collection, "metadata", None) + space = str(raw_meta.get("hnsw:space", "l2")) if isinstance(raw_meta, dict) else "l2" + try: + results = collection.query( + query_embeddings=[vector], # type: ignore[arg-type] # chroma stubs want numpy; lists are accepted + n_results=n, + where={"clearance": {"$in": list(allowed_clearances)}}, # type: ignore[dict-item] + ) + except Exception: + # Kept for parity with the behaviour before this interface existed: ask without the filter and let the + # guard in VectorStore.query() do the filtering. Other adapters fail closed instead. + results = collection.query(query_embeddings=[vector], n_results=max(n, 20)) # type: ignore[arg-type] + docs = (results.get("documents") or [[]])[0] + metas = (results.get("metadatas") or [[]])[0] or [{}] * len(docs) + ids = (results.get("ids") or [[]])[0] + distances = (results.get("distances") or [[]])[0] or [2.0] * len(docs) + return [ + Hit(id=cid, document=doc, metadata=dict(meta or {}), similarity=_similarity(float(distance), space)) + for doc, meta, cid, distance in zip(docs, metas, ids, distances) + ] + + +# --- The factory --------------------------------------------------------------------------------------------------- + +_INSTANCES: dict[tuple, VectorStore] = {} + + +def get_vector_store() -> VectorStore: + """The store selected by VECTOR_STORE, read when called so a settings override applies.""" + kind = settings.vector_store + if kind == "chroma": + return ChromaStore() + if kind == "qdrant": + key: tuple = ("qdrant", settings.qdrant_url, settings.qdrant_api_key, settings.qdrant_prefix) + if key not in _INSTANCES: + from rag_lab.storage.qdrant_store import QdrantStore + + _INSTANCES[key] = QdrantStore(settings.qdrant_url, settings.qdrant_api_key, settings.qdrant_prefix) + return _INSTANCES[key] + if kind == "pgvector": + key = ("pgvector", settings.pgvector_uri, settings.pgvector_schema) + if key not in _INSTANCES: + from rag_lab.storage.pgvector_store import PgvectorStore + + _INSTANCES[key] = PgvectorStore(settings.pgvector_uri, settings.pgvector_schema) + return _INSTANCES[key] + raise VectorStoreError(f"VECTOR_STORE must be one of {', '.join(STORES)}; got '{kind}'.") diff --git a/rag_lab/storage/workspace_manager.py b/rag_lab/storage/workspace_manager.py index 6dacdc4..9e45d86 100644 --- a/rag_lab/storage/workspace_manager.py +++ b/rag_lab/storage/workspace_manager.py @@ -7,20 +7,19 @@ from contextlib import closing from pathlib import Path from typing import List, Dict, Any, Tuple -import chromadb from psycopg_pool import ConnectionPool from rag_lab.config import settings from rag_lab.storage.checkpointer import is_postgres_uri, get_sqlite_path from rag_lab.ingestion.parsers import SUPPORTED_EXTENSIONS from rag_lab.retrieval.sparse_search import delete_sparse_index, has_sparse_index +from rag_lab.storage.vector_store import WorkspaceNotIndexed, get_vector_store logger = logging.getLogger("workspace_manager") # `ws___build` must stay inside Chroma's 63-character collection name limit. _MAX_WORKSPACE_ID = 48 -_TEMP_SUFFIXES = ("__build", "__old") def sanitize_workspace_id(raw_name: str) -> str: @@ -69,15 +68,11 @@ def list_workspaces() -> List[str]: ensure_base_dirs() workspaces = set() - # 1. Discover collections from ChromaDB (not the temporary __build/__old copies used while re-indexing) + # 1. Discover indexes in the vector store (not the temporary copies used while re-indexing) try: - client = chromadb.PersistentClient(path=settings.chroma_path) - for col in client.list_collections(): - name = getattr(col, "name", str(col)) - if name.startswith("ws_") and not name.endswith(_TEMP_SUFFIXES): - workspaces.add(name[3:]) + workspaces |= get_vector_store().workspaces() except Exception as exc: - logger.warning("Could not list Chroma collections: %s", exc) + logger.warning("Could not list the vector store's workspaces: %s", exc) # 2. Discover SQLite FTS5 keyword indices if os.path.exists(settings.bm25_dir): @@ -165,12 +160,11 @@ def create_workspace(raw_name: str, actor: str | None = None) -> str: # Ensure staging folder exists os.makedirs(os.path.join(settings.uploads_dir, ws_id), exist_ok=True) - # Initialize empty collection in ChromaDB + # Initialize an empty index in the vector store try: - client = chromadb.PersistentClient(path=settings.chroma_path) - client.get_or_create_collection(name=f"ws_{ws_id}", metadata={"hnsw:space": "cosine"}) + get_vector_store().ensure(ws_id) except Exception as exc: - logger.warning("Could not create Chroma collection for %s: %s", ws_id, exc) + logger.warning("Could not create the vector index for %s: %s", ws_id, exc) _audit_workspace("workspace_created", ws_id, who) return ws_id @@ -194,16 +188,11 @@ def delete_workspace(raw_name: str, actor: str | None = None) -> bool: if not ws_id: return False - # 1. Delete Chroma collections + # 1. Delete the vector index (and any leftover re-index copy) try: - client = chromadb.PersistentClient(path=settings.chroma_path) - for name in (f"ws_{ws_id}", f"ws_{ws_id}__build", f"ws_{ws_id}__old"): - try: - client.delete_collection(name=name) - except Exception as exc: - logger.debug("Collection %s not deleted: %s", name, exc) + get_vector_store().delete(ws_id) except Exception as exc: - logger.warning("Could not delete Chroma collection for %s: %s", ws_id, exc) + logger.warning("Could not delete the vector index for %s: %s", ws_id, exc) # 2. Delete sparse index (SQLite FTS5 and any leftover legacy PKL) delete_sparse_index(ws_id, base_dir=settings.bm25_dir) @@ -276,13 +265,10 @@ def get_workspace_stats(raw_name: str) -> Dict[str, Any]: bm25_indexed = False embedding = None - # Chunks from Chroma + # Chunks in the vector store try: - from chromadb.errors import NotFoundError - client = chromadb.PersistentClient(path=settings.chroma_path) - col = client.get_collection(name=f"ws_{ws_id}") - chunk_count = col.count() - except NotFoundError: + chunk_count = get_vector_store().count(ws_id) + except WorkspaceNotIndexed: chunk_count = 0 except Exception as exc: logger.warning("Could not read chunk count for %s: %s", ws_id, exc) diff --git a/tests/conftest.py b/tests/conftest.py index ab4cc93..d382c30 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -38,6 +38,16 @@ os.environ["PYTHON_KEYRING_BACKEND"] = "keyring.backends.fail.Keyring" # A .env file may set these for real use; tests get no second anchor location unless they ask for one. os.environ["AUDIT_ANCHOR_COPY_DIR"] = "" +# Scaling options stay off: a .env that enables Qdrant, pgvector or a shared cache must never make a test write to a +# real service. Tests of those options opt in with their own namespace (tests/test_vector_stores.py, test_shared_cache.py). +os.environ["VECTOR_STORE"] = "chroma" +os.environ["CACHE_BACKEND"] = "sqlite" +os.environ["CACHE_DB_URI"] = "" +os.environ["QDRANT_URL"] = "" +os.environ["QDRANT_API_KEY"] = "" +os.environ["PGVECTOR_URI"] = "" +os.environ["LOCAL_LLM_API_KEY"] = "local" +os.environ["LOCAL_LLM_MODEL"] = "" os.environ["BACKUP_DIR"] = os.path.join(_scratch_dir, "backups") os.environ["DB_URI"] = f"sqlite:///{os.path.join(_scratch_dir, 'checkpoints.db')}" # Indexes and uploads too: no test may write to the repository's chroma_data/, bm25_indices/, or uploads/. @@ -58,6 +68,8 @@ raise RuntimeError(f"settings.{_field} was built before the test redirect: {getattr(settings, _field)}") if settings.key_store != "file" or settings.key_envelope or settings.audit_anchor_copy_dir: raise RuntimeError("key storage or the anchor copy directory was built before the test redirect") +if settings.vector_store != "chroma" or settings.cache_backend != "sqlite" or settings.qdrant_url or settings.pgvector_uri: + raise RuntimeError("a scaling option was built before the test redirect: tests must not reach a real service") if not settings.db_uri.startswith(f"sqlite:///{_scratch_dir}"): raise RuntimeError(f"settings.db_uri was built before the test redirect: {settings.db_uri}") diff --git a/tests/test_answers.py b/tests/test_answers.py index 02685d4..bca0f14 100644 --- a/tests/test_answers.py +++ b/tests/test_answers.py @@ -104,7 +104,7 @@ def test_dense_outage_is_recorded_and_sparse_results_still_return(rollback_index def broken_client(*args, **kwargs): raise RuntimeError("vector store offline") - monkeypatch.setattr("rag_lab.pipeline.retrieval.chromadb.PersistentClient", broken_client) + monkeypatch.setattr("chromadb.PersistentClient", broken_client) # what ChromaStore opens state = hybrid_retrieval_node(_retrieval_state(ROLLBACK_QUESTION)) assert "dense_retrieval_unavailable: RuntimeError" in state["degraded"] assert state["documents"], "keyword retrieval should still answer from the index" diff --git a/tests/test_ingest_reuse.py b/tests/test_ingest_reuse.py index 8a9e182..6c60991 100644 --- a/tests/test_ingest_reuse.py +++ b/tests/test_ingest_reuse.py @@ -174,7 +174,7 @@ def test_an_unreadable_previous_index_means_embedding_everything_not_failing(sto first = sync(source) embedder.calls.clear() (source / "alpha.md").write_text("# Alpha\n\nChanged so that a sync is needed at all.") - real, reads = ing.chromadb.PersistentClient, {"n": 0} + real, reads = chromadb.PersistentClient, {"n": 0} class Client: def __init__(self, path): @@ -189,7 +189,7 @@ def get_collection(self, name): def __getattr__(self, item): return getattr(self._client, item) - monkeypatch.setattr(ing.chromadb, "PersistentClient", lambda path: Client(path)) + monkeypatch.setattr(chromadb, "PersistentClient", lambda path: Client(path)) result = sync(source) assert result.reused == 0 and len(embedder.embedded) >= first.chunks - 2 diff --git a/tests/test_rbac.py b/tests/test_rbac.py index 8e28a4d..f44300e 100644 --- a/tests/test_rbac.py +++ b/tests/test_rbac.py @@ -287,7 +287,7 @@ def query_side_effect(**kwargs): fake_collection.query.side_effect = query_side_effect fake_client = mocker.MagicMock() fake_client.get_collection.return_value = fake_collection - mocker.patch("rag_lab.pipeline.retrieval.chromadb.PersistentClient", return_value=fake_client) + mocker.patch("chromadb.PersistentClient", return_value=fake_client) # what ChromaStore opens mocker.patch("rag_lab.pipeline.retrieval.search_sparse_chunks", return_value=[]) out = hybrid_retrieval_node({ diff --git a/tests/test_vector_stores.py b/tests/test_vector_stores.py new file mode 100644 index 0000000..6d3c1aa --- /dev/null +++ b/tests/test_vector_stores.py @@ -0,0 +1,291 @@ +"""Every vector store answers the same questions the same way: ranking, similarity scale, clearance, atomic publish. + +The suite is parametrized over the adapters. Chroma always runs; Qdrant (embedded local mode, and a real server when +QDRANT_TEST_URL is set) and pgvector (a real server, PGVECTOR_TEST_URI) join it when available. REQUIRE_QDRANT=1 / +REQUIRE_PGVECTOR=1 (CI) turn a missing one into a failure instead of a skip. Tests against a real service use their own +namespace (a collection prefix, a schema) and remove it afterwards: never the bare `ws_*` names of a real install. +""" +import os +import uuid + +import numpy as np +import pytest + +from rag_lab.config import settings +from rag_lab.retrieval.embedding_provider import FINGERPRINT_KEY +from rag_lab.storage import vector_store as vs +from rag_lab.storage.vector_store import ( + ChromaStore, + Hit, + VectorStore, + VectorStoreError, + WorkspaceNotIndexed, + get_vector_store, + guard_hits, +) + +DIM = 24 +FINGERPRINT = "fastembed:test/model" +CLEARANCES = ["public", "internal", "confidential", "secret"] + + +def unit_cos(a, b) -> float: + a, b = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64) + return float(np.dot(a, b) / np.linalg.norm(a) / np.linalg.norm(b)) + + +def corpus(workspace: str, n: int = 12, seed: int = 1): + """Chunks with vectors of norm ~5 (like the real embeddings, which are not unit length) and mixed clearances.""" + rng = np.random.default_rng(seed) + vectors = rng.normal(size=(n, DIM)) * 5.0 + ids = [f"{workspace}_{i}" for i in range(n)] + docs = [f"chunk {i} of {workspace}" for i in range(n)] + metas = [{"source": f"doc{i % 3}.md", "workspace": workspace, "clearance": CLEARANCES[i % 4], "id": ids[i]} for i in range(n)] + return ids, docs, metas, vectors + + +def publish(store: VectorStore, workspace: str, n: int = 12, seed: int = 1, fingerprint: str = FINGERPRINT): + ids, docs, metas, vectors = corpus(workspace, n, seed) + store.publish(workspace, ids, docs, metas, vectors.tolist(), fingerprint) + return ids, docs, metas, vectors + + +# --- the adapters under test ----------------------------------------------------------------------------------- + +def _unavailable(reason: str, flag: str): + if os.getenv(flag) == "1": + pytest.fail(f"{flag}=1 but {reason}") + pytest.skip(reason) + + +@pytest.fixture +def chroma(tmp_path, monkeypatch): + monkeypatch.setattr(settings, "chroma_path", str(tmp_path / "chroma")) + return ChromaStore() + + +@pytest.fixture(params=["chroma", "qdrant-local", "qdrant-server", "pgvector"]) +def store(request, tmp_path, monkeypatch): + kind = request.param + if kind == "chroma": + monkeypatch.setattr(settings, "chroma_path", str(tmp_path / "chroma")) + yield ChromaStore() + return + if kind.startswith("qdrant"): + try: + from rag_lab.storage.qdrant_store import QdrantStore + except ImportError: + _unavailable("qdrant-client is not installed", "REQUIRE_QDRANT") + prefix = f"pytest_{uuid.uuid4().hex[:8]}_" + if kind == "qdrant-local": + adapter = QdrantStore(":memory:", "", prefix) + else: + url = os.getenv("QDRANT_TEST_URL") + if not url: + _unavailable("QDRANT_TEST_URL is not set (start one with `make scale-up PROFILE=qdrant`)", "REQUIRE_QDRANT") + adapter = QdrantStore(url, os.getenv("QDRANT_TEST_API_KEY", ""), prefix) + yield adapter + adapter.drop_namespace() + return + uri = os.getenv("PGVECTOR_TEST_URI") + if not uri: + _unavailable("PGVECTOR_TEST_URI is not set (start one with `make scale-up PROFILE=pgvector`)", "REQUIRE_PGVECTOR") + from rag_lab.storage.pgvector_store import PgvectorStore + + schema = f"pytest_{uuid.uuid4().hex[:8]}" + adapter = PgvectorStore(uri, schema) + yield adapter + adapter.drop_namespace() + + +# --- conformance --------------------------------------------------------------------------------------------------- + +def test_an_unindexed_workspace_is_reported_not_guessed(store): + assert not store.exists("ghost") and "ghost" not in store.workspaces() + for call in (lambda: store.count("ghost"), lambda: store.collection_metadata("ghost"), + lambda: store.query("ghost", [0.1] * DIM, 3, CLEARANCES)): + with pytest.raises(WorkspaceNotIndexed): + call() + + +def test_ensure_creates_an_empty_workspace(store): + store.ensure("fresh") + assert store.exists("fresh") and "fresh" in store.workspaces() and store.count("fresh") == 0 + assert store.query("fresh", [0.1] * DIM, 3, CLEARANCES) == [] + store.ensure("fresh") # idempotent + + +def test_publish_stores_chunks_text_metadata_and_the_model_stamp(store): + ids, docs, metas, _ = publish(store, "ops") + assert store.exists("ops") and store.count("ops") == len(ids) + assert store.collection_metadata("ops")[FINGERPRINT_KEY] == FINGERPRINT + hits = store.query("ops", [1.0] * DIM, len(ids), CLEARANCES) + assert {h.id for h in hits} == set(ids) + by_id = {h.id: h for h in hits} + for cid, doc, meta in zip(ids, docs, metas): + assert by_id[cid].document == doc + assert by_id[cid].metadata["source"] == meta["source"] and by_id[cid].metadata["clearance"] == meta["clearance"] + + +def test_similarity_is_cosine_on_chromas_scale_and_ranked_best_first(store): + """The relevance gates are calibrated on `1 - cosine distance`; the value must agree, not just the order.""" + ids, _, metas, vectors = publish(store, "ops") + query = np.random.default_rng(7).normal(size=DIM) * 5.0 + hits = store.query("ops", query.tolist(), 5, CLEARANCES) + expected = sorted(((unit_cos(query, v), cid) for v, cid in zip(vectors, ids)), reverse=True)[:5] + assert [h.id for h in hits] == [cid for _, cid in expected] + for hit, (cosine, _) in zip(hits, expected): + assert abs(hit.similarity - cosine) < 1e-4, (hit.id, hit.similarity, cosine) + assert [h.similarity for h in hits] == sorted((h.similarity for h in hits), reverse=True) + + +def test_the_callers_clearance_limits_what_is_returned(store): + ids, _, metas, vectors = publish(store, "ops") + secret_idx = next(i for i, m in enumerate(metas) if m["clearance"] == "secret") + query = vectors[secret_idx].tolist() # identical to a secret chunk: it would rank first without the filter + public = store.query("ops", query, len(ids), ["public"]) + assert public and all(h.metadata["clearance"] == "public" for h in public) + assert ids[secret_idx] not in {h.id for h in public} + internal = store.query("ops", query, len(ids), ["public", "internal"]) + assert {h.metadata["clearance"] for h in internal} <= {"public", "internal"} and len(internal) > len(public) + everyone = store.query("ops", query, 3, CLEARANCES) + assert everyone[0].id == ids[secret_idx] and everyone[0].similarity > 0.9999 + + +def test_a_workspace_never_returns_another_workspaces_chunks(store): + """Same vectors in both: a shared table (pgvector) or a shared prefix must still keep them apart.""" + a_ids, *_ = publish(store, "alpha", seed=3) + b_ids, *_ = publish(store, "beta", seed=3) + hits = store.query("alpha", [1.0] * DIM, 50, CLEARANCES) + assert hits and {h.id for h in hits} <= set(a_ids) and not ({h.id for h in hits} & set(b_ids)) + assert store.workspaces() >= {"alpha", "beta"} + + +def test_publish_replaces_the_whole_index(store): + old_ids, *_ = publish(store, "ops", n=12, seed=1) + new_ids, *_ = publish(store, "ops", n=5, seed=2, fingerprint="fastembed:another/model") + assert store.count("ops") == 5 + assert {h.id for h in store.query("ops", [1.0] * DIM, 50, CLEARANCES)} == set(new_ids) + assert set(old_ids) - set(new_ids), "chunks that only the old index had are gone" + assert store.collection_metadata("ops")[FINGERPRINT_KEY] == "fastembed:another/model" + + +def test_a_publish_that_fails_leaves_the_old_index_serving(store): + old_ids, *_ = publish(store, "ops", n=8, seed=1) + ids, docs, metas, vectors = corpus("ops", 6, seed=9) + with pytest.raises(Exception): + store.publish("ops", ids, docs, metas, vectors[:3].tolist(), FINGERPRINT) # 6 chunks, 3 vectors + with pytest.raises(Exception): + store.publish("ops", ids[:2], docs[:2], metas[:2], [[0.1] * DIM, [0.1] * (DIM + 5)], FINGERPRINT) # mixed sizes + assert store.count("ops") == 8 + assert {h.id for h in store.query("ops", [1.0] * DIM, 50, CLEARANCES)} == set(old_ids) + assert store.collection_metadata("ops")[FINGERPRINT_KEY] == FINGERPRINT + assert not [w for w in store.workspaces() if "__" in w], "no half-built copy is left behind" + + +def test_read_vectors_returns_each_chunk_text_with_a_vector_pointing_the_same_way(store): + ids, docs, _, vectors = publish(store, "ops") + got = dict(store.read_vectors("ops")) + assert set(got) == set(docs) + for doc, vector in zip(docs, vectors): + assert unit_cos(got[doc], vector) > 0.99999 # a store may normalise; the direction is what ranking uses + with pytest.raises(WorkspaceNotIndexed): + list(store.read_vectors("ghost")) + + +def test_delete_removes_the_workspace_and_is_quiet_when_there_is_none(store): + publish(store, "ops") + store.delete("ops") + assert not store.exists("ops") and "ops" not in store.workspaces() + store.delete("ops") + store.delete("never-existed") + + +def test_a_reused_vector_ranks_like_a_fresh_one(store): + """The ingest path re-publishes vectors it read back: ranking must not change.""" + ids, docs, metas, vectors = publish(store, "ops") + before = [h.id for h in store.query("ops", [1.0] * DIM, 6, CLEARANCES)] + reread = dict(store.read_vectors("ops")) + store.publish("ops", ids, docs, metas, [reread[d] for d in docs], FINGERPRINT) + assert [h.id for h in store.query("ops", [1.0] * DIM, 6, CLEARANCES)] == before + + +# --- the guard belongs to the base class, not to the adapters ------------------------------------------------------------ + +class LeakyStore(VectorStore): + """An adapter whose server-side filter is broken: it returns every chunk it holds, whatever was asked.""" + + name = "leaky" + + def __init__(self, hits): + self._hits = hits + + def workspaces(self): return {"ops"} + def ensure(self, workspace): pass + def exists(self, workspace): return True + def count(self, workspace): return len(self._hits) + def collection_metadata(self, workspace): return {FINGERPRINT_KEY: FINGERPRINT} + def read_vectors(self, workspace): return iter(()) + def publish(self, *a, **k): pass + def delete(self, workspace): pass + + def _search(self, workspace, vector, n, allowed_clearances): + return list(self._hits) + + +def _hit(cid, clearance, workspace="ops", similarity=0.99): + meta = {"source": "x.md", "workspace": workspace} + if clearance is not None: + meta["clearance"] = clearance + return Hit(cid, f"text {cid}", meta, similarity) + + +def test_the_base_class_guard_stops_a_chunk_even_when_the_adapters_filter_leaks(): + leaky = LeakyStore([ + _hit("pub", "public"), _hit("sec", "secret"), _hit("conf", "Confidential "), _hit("none", None), + _hit("blank", ""), _hit("junk", "top-secret"), _hit("other-ws", "public", workspace="hr"), + _hit("upper", "PUBLIC"), + ]) + assert [h.id for h in leaky.query("ops", [0.0], 10, ["public"])] == ["pub", "upper"] + assert [h.id for h in leaky.query("ops", [0.0], 10, ["public", "internal", "confidential"])] == ["pub", "conf", "upper"] + assert leaky._search("ops", [0.0], 10, ["public"]) and len(leaky._search("ops", [0.0], 10, ["public"])) == 8, \ + "the leak is real: without the guard all eight come back" + + +def test_the_guard_is_applied_by_query_not_left_to_each_adapter(monkeypatch): + leaky = LeakyStore([_hit("sec", "secret")]) + monkeypatch.setattr(vs, "guard_hits", lambda hits, ws, allowed: list(hits)) # mutate: no guard + assert [h.id for h in leaky.query("ops", [0.0], 5, ["public"])] == ["sec"], "with the guard removed it leaks" + + +def test_guard_hits_in_isolation(): + hits = [_hit("a", "public"), _hit("b", "secret")] + assert guard_hits(hits, "ops", ["public"]) == [hits[0]] + assert guard_hits(hits, "ops", []) == [] + + +def test_chroma_keeps_its_unfiltered_fallback_for_old_collections_and_the_guard_still_applies(chroma): + """Parity with the behaviour before the interface existed (tests/test_rbac.py exercises it end to end).""" + publish(chroma, "ops") + from unittest import mock + + collection = chroma._collection("ops") + real = collection.query + + def flaky(**kw): + if "where" in kw: + raise RuntimeError("legacy collection without a where index") + return real(**kw) + + with mock.patch.object(ChromaStore, "_collection", return_value=mock.Mock(query=flaky, metadata=collection.metadata)): + hits = chroma.query("ops", [1.0] * DIM, 5, ["public"]) + assert hits and all(h.metadata["clearance"] == "public" for h in hits) + + +# --- the factory ------------------------------------------------------------------------------------------------------------ + +def test_the_default_store_is_embedded_chroma_and_an_unknown_one_is_refused(monkeypatch): + assert settings.vector_store == "chroma" and isinstance(get_vector_store(), ChromaStore) + monkeypatch.setattr(settings, "vector_store", "milvus") + with pytest.raises(VectorStoreError, match="chroma, qdrant, pgvector"): + get_vector_store() From dcd7fbdff002c10e4009d5f9eede218d394b5aec Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 20:31:55 +0200 Subject: [PATCH 2/8] feat(storage): optional Qdrant and pgvector vector stores, and their compose profiles VECTOR_STORE=qdrant|pgvector (default chroma, unchanged). Qdrant: a versioned collection behind an alias that moves in one atomic request; the model stamp is the collection's metadata. pgvector: one transaction per publish over rag_chunks/rag_workspaces in a configurable schema; no new Python package. Neither has an unfiltered fallback: a failing filter makes dense retrieval unavailable, labelled. The clearance guard stays in VectorStore.query() for both. The shared suite runs every adapter (Chroma, embedded Qdrant, a real Qdrant 1.20 server, a real pgvector server): similarity values within 1e-4 of Chroma's scale and of numpy's cosine, identical ranking against Chroma on the same corpus, server-side clearance filtering, workspace isolation in pgvector's shared table, and a failed publish leaving the old index serving. make eval passes through each real adapter with metrics identical to Chroma's (answer rate and recall 0.9773, abstention 1.0). Mutations (no server-side filter; the wrong score scale) fail the suite. qdrant-client is an optional extra (five small packages). Compose profiles qdrant, pgvector and vllm (loopback; make scale-up SERVICE=...); make up still starts only Postgres. The pgvector pool is closed at exit (Python 3.14 raises otherwise). --- Makefile | 16 ++- docker-compose.yml | 67 ++++++++++ pyproject.toml | 4 + rag_lab/storage/pgvector_store.py | 185 ++++++++++++++++++++++++++ rag_lab/storage/qdrant_store.py | 208 ++++++++++++++++++++++++++++++ rag_lab/storage/vector_store.py | 9 +- tests/test_vector_stores.py | 94 +++++++++++++- 7 files changed, 575 insertions(+), 8 deletions(-) create mode 100644 rag_lab/storage/pgvector_store.py create mode 100644 rag_lab/storage/qdrant_store.py diff --git a/Makefile b/Makefile index 5f45694..63f9428 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ .DEFAULT_GOAL := help -.PHONY: help status setup rust-build models verify-models calibrate-models up down ollama-pull test test-offline lint typecheck coverage workspaces create-workspace delete-workspace sync ingest ui audit-verify audit-tail audit-feedback audit-anchor audit-export retention backup backups backup-verify add-user grant revoke set-clearance set-admin set-role set-default team-mode groups create-group delete-group add-to-group remove-from-group grant-group set-password revoke-sessions delete-user users reset-users factory-reset recover db-password lan-tls lan-tls-down lan-tls-ca migrate-keys totp-enroll totp-recovery totp-disable eval eval-judge eval-llm eval-calibrate graph clean reviews review approve reject +.PHONY: help status setup rust-build models verify-models calibrate-models up down ollama-pull test test-offline lint typecheck coverage workspaces create-workspace delete-workspace sync ingest ui audit-verify audit-tail audit-feedback audit-anchor audit-export retention backup backups backup-verify add-user grant revoke set-clearance set-admin set-role set-default team-mode groups create-group delete-group add-to-group remove-from-group grant-group set-password revoke-sessions delete-user users reset-users factory-reset recover db-password scale-up scale-down lan-tls lan-tls-down lan-tls-ca migrate-keys totp-enroll totp-recovery totp-disable eval eval-judge eval-llm eval-calibrate graph clean reviews review approve reject # 0. Command Reference & System Readiness help: @@ -13,6 +13,8 @@ help: @echo " make rust-build Compile the native Rust accelerator (rag_accel) from source" @echo " make ollama-pull Pull the optional Ollama chat model" @echo " make up / make down Start / stop the optional PostgreSQL checkpointer (Docker); 'up' creates .env with a random password" + @echo " make scale-up SERVICE= Start an optional scaling service (docs/scaling.md); make up never does" + @echo " make scale-down SERVICE= Stop it (its container and data volume are kept)" @echo " make lan-tls Optional HTTPS for the UI on the LAN (Caddy profile, docs/lan-tls.md)" @echo " make lan-tls-down Stop the LAN TLS proxy" @echo " make lan-tls-ca Export Caddy's local root certificate for client devices" @@ -129,6 +131,18 @@ up: @$(if $(wildcard venv/bin/python),./venv/bin/python,python3) -m rag_lab.ops.init_env docker compose up -d +# Optional scaling services (compose profiles qdrant, pgvector, vllm). Nothing here is needed on one laptop. The +# service only starts: turn it on in .env (VECTOR_STORE=qdrant, VECTOR_STORE=pgvector, LLM_PROVIDER=openai). +scale-up: + @case "$(SERVICE)" in qdrant|pgvector|vllm) ;; *) echo "Usage: make scale-up SERVICE=qdrant|pgvector|vllm"; exit 1;; esac + @docker info >/dev/null 2>&1 || { echo "Error: Docker daemon is not running."; exit 1; } + @$(if $(wildcard venv/bin/python),./venv/bin/python,python3) -m rag_lab.ops.init_env + docker compose --profile $(SERVICE) up -d $(SERVICE) + +scale-down: + @case "$(SERVICE)" in qdrant|pgvector|vllm) ;; *) echo "Usage: make scale-down SERVICE=qdrant|pgvector|vllm"; exit 1;; esac + docker compose --profile $(SERVICE) stop $(SERVICE) + # Optional LAN TLS (compose profile "caddy"): Caddy terminates HTTPS with its own local CA and proxies to the # Streamlit UI on the host. Never started by `make up`. Set RAG_LAN_HOST and RAG_HTTPS_PORT in .env. lan-tls: diff --git a/docker-compose.yml b/docker-compose.yml index 9a2b571..f3c841a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -40,7 +40,74 @@ services: - caddy_data:/data - caddy_config:/config + # --- Scaling options. Each is a profile: 'make up' starts none of them, 'make scale-up SERVICE=' starts one. + # They bind to loopback; put them behind a TLS proxy or a private network before reaching them from another host. + # Setup, trade-offs and what each does not solve: docs/scaling.md. + + # Vector store: VECTOR_STORE=qdrant. An empty QDRANT_API_KEY means no authentication (loopback only); set one + # before exposing the port. + qdrant: + profiles: ["qdrant"] + image: qdrant/qdrant:v1.20.0 + container_name: local-rag-qdrant + restart: unless-stopped + environment: + QDRANT__SERVICE__API_KEY: ${QDRANT_API_KEY:-} + ports: + - "127.0.0.1:${QDRANT_PORT:-6333}:6333" + volumes: + - qdrant_data:/qdrant/storage + + # Vector store: VECTOR_STORE=pgvector, and the shared cache (CACHE_BACKEND=postgres) if you want one database. This + # is a separate server from the checkpoint database because postgres:16-alpine has no pgvector extension. + pgvector: + profiles: ["pgvector"] + image: pgvector/pgvector:pg16 + container_name: local-rag-pgvector + restart: unless-stopped + environment: + POSTGRES_DB: rag_vectors + POSTGRES_USER: ${POSTGRES_USER:-rag_user} + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?Set POSTGRES_PASSWORD in .env (run 'make up', which creates .env)} + ports: + - "127.0.0.1:${PGVECTOR_PORT:-5433}:5432" + volumes: + - pgvector_data:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER} -d rag_vectors"] + interval: 5s + timeout: 5s + retries: 5 + + # Chat model: LLM_PROVIDER=openai LOCAL_LLM_URL=http://localhost:8000/v1 LOCAL_LLM_MODEL=. Needs an NVIDIA + # GPU and the NVIDIA Container Toolkit (it cannot run on a Mac). The first start downloads the model from Hugging Face: + # unlike the ONNX models, it is not SHA-256 pinned, and the host is not air-gapped until the cache volume is filled. + vllm: + profiles: ["vllm"] + image: vllm/vllm-openai:latest + container_name: local-rag-vllm + restart: unless-stopped + command: ["--model", "${VLLM_MODEL:-Qwen/Qwen3-8B}", "--host", "0.0.0.0", "--port", "8000"] + environment: + VLLM_API_KEY: ${VLLM_API_KEY:-} + HF_TOKEN: ${HF_TOKEN:-} + ports: + - "127.0.0.1:${VLLM_PORT:-8000}:8000" + volumes: + - vllm_cache:/root/.cache/huggingface + ipc: host + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + volumes: pgdata: caddy_data: caddy_config: + qdrant_data: + pgvector_data: + vllm_cache: diff --git a/pyproject.toml b/pyproject.toml index 5a8729d..94acf5f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,10 @@ dev = [ rust = [ "maturin>=1.7.0", ] +# Optional scaling options (docs/scaling.md). pgvector needs no extra: psycopg is already a dependency. +qdrant = [ + "qdrant-client>=1.9", +] [project.scripts] local-rag = "rag_lab.storage.workspace_manager:main" diff --git a/rag_lab/storage/pgvector_store.py b/rag_lab/storage/pgvector_store.py new file mode 100644 index 0000000..68b3eaa --- /dev/null +++ b/rag_lab/storage/pgvector_store.py @@ -0,0 +1,185 @@ +"""PostgreSQL with the pgvector extension as the vector store (VECTOR_STORE=pgvector). Needs no new Python package. + +Layout, in one schema (PGVECTOR_SCHEMA, `public` by default): + + rag_workspaces(workspace PRIMARY KEY, fingerprint, created_at) -- one row per workspace, even an empty one + rag_chunks(workspace, id, document, metadata jsonb, clearance, embedding vector) -- PRIMARY KEY (workspace, id) + +All workspaces share the one table, which is why every query filters on `workspace` and why VectorStore.query() checks +the workspace of each hit again. A publish is one transaction (delete the workspace's rows, insert the new ones, update +its fingerprint): readers see the old index until the commit, and a failure rolls back to it. + +Scores. `<=>` is pgvector's cosine *distance*, so `1 - (embedding <=> query)` is the cosine similarity, the same scale as +Chroma's `1 - distance`. + +No ANN index is created. The column has no fixed dimension (the model can change), HNSW needs one, and an exact scan of +a workspace's rows is what Chroma's results are compared with. For workspaces beyond a few hundred thousand chunks add an +HNSW index on the cast column yourself (docs/scaling.md). There is no unfiltered fallback: a failing filter makes dense +retrieval unavailable (labelled) instead of asking for everything. +""" +from __future__ import annotations + +import atexit +import json +import re +from collections.abc import Iterator, Sequence +from typing import Any + +from psycopg import sql +from psycopg.types.json import Jsonb +from psycopg_pool import ConnectionPool + +from rag_lab.config import settings +from rag_lab.retrieval.embedding_provider import FINGERPRINT_KEY +from rag_lab.storage.vector_store import Hit, VectorStore, VectorStoreError, WorkspaceNotIndexed + +_SCHEMA = re.compile(r"[a-z_][a-z0-9_]{0,62}") +_BATCH = 500 + + +def _literal(vector: Sequence[float]) -> str: + """pgvector's text form of a vector: '[0.1,0.2,...]'. Passed as a parameter and cast with ::vector.""" + return "[" + ",".join(repr(float(x)) for x in vector) + "]" + + +class PgvectorStore(VectorStore): + name = "pgvector" + + def __init__(self, uri: str, schema: str = "public") -> None: + if not uri: + raise VectorStoreError( + "VECTOR_STORE=pgvector needs PGVECTOR_URI, or POSTGRES_PASSWORD in .env for the compose profile " + "(`make scale-up PROFILE=pgvector`)." + ) + if not _SCHEMA.fullmatch(schema): + raise VectorStoreError(f"PGVECTOR_SCHEMA must be a plain lowercase identifier; got '{schema}'.") + self._schema = schema + self._pool = ConnectionPool(conninfo=uri, min_size=1, max_size=max(2, settings.checkpoint_pool_size), open=True, + timeout=15, kwargs={"autocommit": False}) + # The pool's worker threads must be stopped before the interpreter finalises (Python 3.14 raises otherwise). + atexit.register(self.close) + try: + self._pool.wait(timeout=15) + self._create_tables() + except Exception as exc: + self._pool.close() + raise VectorStoreError( + f"Could not prepare pgvector in {uri.split('@')[-1]} ({type(exc).__name__}: {exc}). The server needs the " + "pgvector extension: use the `pgvector` compose profile (image pgvector/pgvector), not plain postgres." + ) from exc + + def close(self) -> None: + self._pool.close() + + # --- SQL helpers ---------------------------------------------------------------------------------------------- + def _t(self, table: str) -> sql.Composed: + return sql.SQL("{}.{}").format(sql.Identifier(self._schema), sql.Identifier(table)) + + def _create_tables(self) -> None: + with self._pool.connection() as conn: + conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + conn.execute(sql.SQL("CREATE SCHEMA IF NOT EXISTS {}").format(sql.Identifier(self._schema))) + conn.execute(sql.SQL( + "CREATE TABLE IF NOT EXISTS {} (workspace text PRIMARY KEY, fingerprint text, " + "created_at timestamptz NOT NULL DEFAULT now())").format(self._t("rag_workspaces"))) + conn.execute(sql.SQL( + "CREATE TABLE IF NOT EXISTS {} (workspace text NOT NULL, id text NOT NULL, document text NOT NULL, " + "metadata jsonb NOT NULL, clearance text NOT NULL, embedding vector NOT NULL, " + "PRIMARY KEY (workspace, id))").format(self._t("rag_chunks"))) + conn.execute(sql.SQL("CREATE INDEX IF NOT EXISTS {} ON {} (workspace, clearance)").format( + sql.Identifier(f"rag_chunks_clearance_{self._schema}"[:63]), self._t("rag_chunks"))) + + def _scalar(self, query: sql.Composed, params: tuple) -> Any: + with self._pool.connection() as conn: + row = conn.execute(query, params).fetchone() + return row[0] if row else None + + # --- the interface ------------------------------------------------------------------------------------------------ + def workspaces(self) -> set[str]: + with self._pool.connection() as conn: + rows = conn.execute(sql.SQL("SELECT workspace FROM {}").format(self._t("rag_workspaces"))).fetchall() + return {r[0] for r in rows} + + def exists(self, workspace: str) -> bool: + return bool(self._scalar(sql.SQL("SELECT 1 FROM {} WHERE workspace = %s").format(self._t("rag_workspaces")), (workspace,)) == 1) + + def _require(self, workspace: str) -> None: + if not self.exists(workspace): + raise WorkspaceNotIndexed(workspace) + + def ensure(self, workspace: str) -> None: + with self._pool.connection() as conn: + conn.execute(sql.SQL("INSERT INTO {} (workspace) VALUES (%s) ON CONFLICT DO NOTHING").format( + self._t("rag_workspaces")), (workspace,)) + + def count(self, workspace: str) -> int: + self._require(workspace) + return int(self._scalar(sql.SQL("SELECT count(*) FROM {} WHERE workspace = %s").format(self._t("rag_chunks")), (workspace,))) + + def collection_metadata(self, workspace: str) -> dict[str, Any] | None: + self._require(workspace) + fingerprint = self._scalar(sql.SQL("SELECT fingerprint FROM {} WHERE workspace = %s").format(self._t("rag_workspaces")), (workspace,)) + return {FINGERPRINT_KEY: fingerprint} if fingerprint else {} + + def read_vectors(self, workspace: str) -> Iterator[tuple[str, list[float]]]: + self._require(workspace) + with self._pool.connection() as conn: + rows = conn.execute( + sql.SQL("SELECT document, embedding::real[] FROM {} WHERE workspace = %s ORDER BY id").format(self._t("rag_chunks")), + (workspace,), + ).fetchall() + for document, vector in rows: + yield str(document), [float(x) for x in vector] + + def publish(self, workspace, ids, documents, metadatas, vectors, fingerprint) -> None: + # Checked before the transaction opens: pgvector would store vectors of mixed sizes in an unsized column. + if not (len(ids) == len(documents) == len(metadatas) == len(vectors)): + raise ValueError("ids, documents, metadatas and vectors must be the same length") + if not ids: + raise ValueError("nothing to publish") + sizes = {len(v) for v in vectors} + if len(sizes) != 1 or 0 in sizes: + raise ValueError(f"vectors must all have the same non-zero size; got {sorted(sizes)}") + rows = [ + (workspace, cid, doc, Jsonb({**meta, "id": cid, "workspace": workspace}), str(meta.get("clearance", "")).strip().lower(), _literal(vec)) + for cid, doc, meta, vec in zip(ids, documents, metadatas, vectors) + ] + with self._pool.connection() as conn: # one transaction: commit on success, roll back on any error + conn.execute( + sql.SQL("INSERT INTO {} (workspace, fingerprint) VALUES (%s, %s) " + "ON CONFLICT (workspace) DO UPDATE SET fingerprint = EXCLUDED.fingerprint").format(self._t("rag_workspaces")), + (workspace, fingerprint), + ) + conn.execute(sql.SQL("DELETE FROM {} WHERE workspace = %s").format(self._t("rag_chunks")), (workspace,)) + insert = sql.SQL("INSERT INTO {} (workspace, id, document, metadata, clearance, embedding) " + "VALUES (%s, %s, %s, %s, %s, %s::vector)").format(self._t("rag_chunks")) + with conn.cursor() as cursor: + for start in range(0, len(rows), _BATCH): + cursor.executemany(insert, rows[start:start + _BATCH]) + + def delete(self, workspace: str) -> None: + with self._pool.connection() as conn: + conn.execute(sql.SQL("DELETE FROM {} WHERE workspace = %s").format(self._t("rag_chunks")), (workspace,)) + conn.execute(sql.SQL("DELETE FROM {} WHERE workspace = %s").format(self._t("rag_workspaces")), (workspace,)) + + def _search(self, workspace: str, vector: list[float], n: int, allowed_clearances: Sequence[str]) -> list[Hit]: + self._require(workspace) + literal = _literal(vector) + with self._pool.connection() as conn: + rows = conn.execute( + sql.SQL("SELECT id, document, metadata, clearance, 1 - (embedding <=> %s::vector) AS similarity FROM {} " + "WHERE workspace = %s AND clearance = ANY(%s) ORDER BY embedding <=> %s::vector LIMIT %s").format(self._t("rag_chunks")), + (literal, workspace, [c.strip().lower() for c in allowed_clearances], literal, n), + ).fetchall() + hits = [] + for cid, document, metadata, clearance, similarity in rows: + meta = dict(metadata if isinstance(metadata, dict) else json.loads(metadata)) + meta["clearance"], meta["workspace"] = clearance, workspace # the filtered columns are the truth + hits.append(Hit(id=str(cid), document=str(document), metadata=meta, similarity=float(similarity))) + return hits + + def drop_namespace(self) -> None: + """Remove this schema and everything in it. For tests and for retiring an install's namespace.""" + with self._pool.connection() as conn: + conn.execute(sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(sql.Identifier(self._schema))) + self._pool.close() diff --git a/rag_lab/storage/qdrant_store.py b/rag_lab/storage/qdrant_store.py new file mode 100644 index 0000000..b214f8f --- /dev/null +++ b/rag_lab/storage/qdrant_store.py @@ -0,0 +1,208 @@ +"""Qdrant as the vector store (VECTOR_STORE=qdrant). Optional: needs the `qdrant-client` package. + +Layout. Each workspace has a *physical* collection `ws___` and an *alias* `ws_` that +points to the live one. Queries go through the alias. A publish builds a new physical collection beside the live one +and then moves the alias in a single request (delete + create is one atomic operation on the server), so readers see +the old index or the new one, and a publish that fails part-way never touches the alias. The model stamp (fingerprint) +is the collection's own metadata. + +Scores. Qdrant's Cosine distance returns the cosine similarity as the score, the same scale as Chroma's +`1 - distance` that the relevance gates are calibrated on (tests/test_vector_stores.py asserts the values). + +Differences from Chroma that matter: vectors are normalised on insert (direction is kept, which is all ranking uses), a +collection's vector size is fixed when it is created (so an empty workspace holds a one-dimensional placeholder until the +first publish), and there is no unfiltered fallback: a failing server-side filter makes dense retrieval unavailable +(labelled) instead of asking for everything and trusting the guard. +""" +from __future__ import annotations + +import logging +import uuid +from collections.abc import Iterator, Sequence +from typing import Any, cast + +from qdrant_client import QdrantClient, models + +from rag_lab.retrieval.embedding_provider import FINGERPRINT_KEY +from rag_lab.storage.vector_store import Hit, VectorStore, VectorStoreError, WorkspaceNotIndexed + +logger = logging.getLogger("qdrant_store") + +_BATCH = 256 +_PLACEHOLDER = models.VectorParams(size=1, distance=models.Distance.COSINE) + + +class QdrantStore(VectorStore): + name = "qdrant" + + def __init__(self, url: str, api_key: str = "", prefix: str = "") -> None: + if not url: + raise VectorStoreError("VECTOR_STORE=qdrant needs QDRANT_URL (the compose profile listens on http://localhost:6333).") + self._prefix = prefix + self._client = QdrantClient(location=":memory:") if url == ":memory:" else QdrantClient( + url=url, api_key=api_key or None, timeout=30 + ) + + # --- naming ------------------------------------------------------------------------------------------------ + def _alias(self, workspace: str) -> str: + return f"{self._prefix}ws_{workspace}" + + def _aliases(self) -> dict[str, str]: + """{alias: physical collection} for this install's workspaces only.""" + head = f"{self._prefix}ws_" + return {a.alias_name: a.collection_name for a in self._client.get_aliases().aliases if a.alias_name.startswith(head)} + + def _live(self, workspace: str) -> str | None: + return self._aliases().get(self._alias(workspace)) + + def _require(self, workspace: str) -> str: + live = self._live(workspace) + if live is None: + raise WorkspaceNotIndexed(workspace) + return live + + # --- the interface ------------------------------------------------------------------------------------------- + def workspaces(self) -> set[str]: + head = f"{self._prefix}ws_" + return {alias[len(head):] for alias in self._aliases()} + + def exists(self, workspace: str) -> bool: + return self._live(workspace) is not None + + def ensure(self, workspace: str) -> None: + if self._live(workspace) is not None: + return + physical = f"{self._alias(workspace)}__{uuid.uuid4().hex[:12]}" + self._client.create_collection(physical, vectors_config=_PLACEHOLDER, metadata={}) + self._set_alias(self._alias(workspace), physical, previous=None) + + def count(self, workspace: str) -> int: + self._require(workspace) + return int(self._client.count(self._alias(workspace), exact=True).count) + + def collection_metadata(self, workspace: str) -> dict[str, Any] | None: + info = self._client.get_collection(self._require(workspace)) + return dict(getattr(info.config, "metadata", None) or {}) + + def read_vectors(self, workspace: str) -> Iterator[tuple[str, list[float]]]: + self._require(workspace) + offset = None + while True: + points, offset = self._client.scroll( + self._alias(workspace), limit=500, offset=offset, with_payload=["document"], with_vectors=True + ) + for point in points: + vector = point.vector + if not isinstance(vector, list) or not point.payload or "document" not in point.payload: + raise VectorStoreError(f"the index of '{workspace}' returned a chunk without its text or vector") + yield str(point.payload["document"]), [float(x) for x in cast(list[float], vector)] + if offset is None: + return + + def publish(self, workspace, ids, documents, metadatas, vectors, fingerprint) -> None: + if not (len(ids) == len(documents) == len(metadatas) == len(vectors)): + raise ValueError("ids, documents, metadatas and vectors must be the same length") + if not ids: + raise ValueError("nothing to publish") + sizes = {len(v) for v in vectors} + if len(sizes) != 1 or 0 in sizes: + raise ValueError(f"vectors must all have the same non-zero size; got {sorted(sizes)}") + alias = self._alias(workspace) + previous = self._aliases().get(alias) + physical = f"{alias}__{uuid.uuid4().hex[:12]}" + self._client.create_collection( + physical, + vectors_config=models.VectorParams(size=sizes.pop(), distance=models.Distance.COSINE), + metadata={FINGERPRINT_KEY: fingerprint}, + ) + try: + try: + self._client.create_payload_index(physical, "clearance", models.PayloadSchemaType.KEYWORD) + except Exception: # the embedded mode has no payload indexes; the server filters correctly without one + logger.debug("No payload index on clearance for %s", physical) + for start in range(0, len(ids), _BATCH): + end = start + _BATCH + self._client.upsert( + physical, + points=[ + models.PointStruct( + id=str(uuid.uuid5(uuid.NAMESPACE_URL, cid)), + vector=[float(x) for x in vector], + payload={**meta, "id": cid, "document": doc, "workspace": workspace}, + ) + for cid, doc, meta, vector in zip(ids[start:end], documents[start:end], metadatas[start:end], vectors[start:end]) + ], + wait=True, + ) + self._set_alias(alias, physical, previous) + except BaseException: + self._drop_collection(physical) + raise + if previous and previous != physical: + self._drop_collection(previous) + + def delete(self, workspace: str) -> None: + alias = self._alias(workspace) + live = self._aliases().get(alias) + if live is not None: + try: + self._client.update_collection_aliases( + change_aliases_operations=[models.DeleteAliasOperation(delete_alias=models.DeleteAlias(alias_name=alias))] + ) + except Exception as exc: + logger.debug("Alias %s was not removed (%s)", alias, exc) + for collection in self._client.get_collections().collections: + if collection.name.startswith(f"{alias}__"): + self._drop_collection(collection.name) + + def _search(self, workspace: str, vector: list[float], n: int, allowed_clearances: Sequence[str]) -> list[Hit]: + self._require(workspace) + if self.count(workspace) == 0: + return [] # an empty workspace holds a one-dimensional placeholder: nothing to compare with + response = self._client.query_points( + self._alias(workspace), + query=[float(x) for x in vector], + limit=n, + with_payload=True, + query_filter=models.Filter( + must=[models.FieldCondition(key="clearance", match=models.MatchAny(any=list(allowed_clearances)))] + ), + ) + hits = [] + for point in response.points: + payload = dict(point.payload or {}) + hits.append(Hit( + id=str(payload.get("id", point.id)), + document=str(payload.pop("document", "")), + metadata=payload, + similarity=float(point.score), + )) + return hits + + # --- helpers ----------------------------------------------------------------------------------------------- + def _set_alias(self, alias: str, physical: str, previous: str | None) -> None: + """Point the alias at `physical`, moving it from `previous` in one atomic request.""" + operations: list[Any] = [] + if previous is not None: + operations.append(models.DeleteAliasOperation(delete_alias=models.DeleteAlias(alias_name=alias))) + operations.append(models.CreateAliasOperation(create_alias=models.CreateAlias(collection_name=physical, alias_name=alias))) + self._client.update_collection_aliases(change_aliases_operations=operations) + + def _drop_collection(self, name: str) -> None: + try: + self._client.delete_collection(name) + except Exception as exc: + logger.debug("Collection %s was not deleted (%s)", name, exc) + + def drop_namespace(self) -> None: + """Remove every alias and collection of this prefix. For tests and for retiring an install's namespace.""" + for alias in list(self._aliases()): + try: + self._client.update_collection_aliases( + change_aliases_operations=[models.DeleteAliasOperation(delete_alias=models.DeleteAlias(alias_name=alias))] + ) + except Exception as exc: + logger.debug("Alias %s was not removed (%s)", alias, exc) + for collection in self._client.get_collections().collections: + if collection.name.startswith(f"{self._prefix}ws_"): + self._drop_collection(collection.name) diff --git a/rag_lab/storage/vector_store.py b/rag_lab/storage/vector_store.py index f784733..e77a21a 100644 --- a/rag_lab/storage/vector_store.py +++ b/rag_lab/storage/vector_store.py @@ -268,8 +268,13 @@ def get_vector_store() -> VectorStore: if kind == "qdrant": key: tuple = ("qdrant", settings.qdrant_url, settings.qdrant_api_key, settings.qdrant_prefix) if key not in _INSTANCES: - from rag_lab.storage.qdrant_store import QdrantStore - + try: + from rag_lab.storage.qdrant_store import QdrantStore + except ImportError as exc: + raise VectorStoreError( + "VECTOR_STORE=qdrant needs the optional qdrant-client package: pip install 'qdrant-client>=1.9' " + "(or pip install -e '.[qdrant]')." + ) from exc _INSTANCES[key] = QdrantStore(settings.qdrant_url, settings.qdrant_api_key, settings.qdrant_prefix) return _INSTANCES[key] if kind == "pgvector": diff --git a/tests/test_vector_stores.py b/tests/test_vector_stores.py index 6d3c1aa..eab2052 100644 --- a/tests/test_vector_stores.py +++ b/tests/test_vector_stores.py @@ -82,18 +82,27 @@ def store(request, tmp_path, monkeypatch): else: url = os.getenv("QDRANT_TEST_URL") if not url: - _unavailable("QDRANT_TEST_URL is not set (start one with `make scale-up PROFILE=qdrant`)", "REQUIRE_QDRANT") + _unavailable("QDRANT_TEST_URL is not set (start one with `make scale-up SERVICE=qdrant`)", "REQUIRE_QDRANT") adapter = QdrantStore(url, os.getenv("QDRANT_TEST_API_KEY", ""), prefix) + try: + adapter._client.get_collections() + except Exception as exc: + _unavailable(f"Qdrant is not reachable at {url} ({type(exc).__name__})", "REQUIRE_QDRANT") yield adapter adapter.drop_namespace() return - uri = os.getenv("PGVECTOR_TEST_URI") - if not uri: - _unavailable("PGVECTOR_TEST_URI is not set (start one with `make scale-up PROFILE=pgvector`)", "REQUIRE_PGVECTOR") + from rag_lab.config import pgvector_uri_from_env from rag_lab.storage.pgvector_store import PgvectorStore + uri = os.getenv("PGVECTOR_TEST_URI") or pgvector_uri_from_env() + if not uri: + _unavailable("no pgvector server: set PGVECTOR_TEST_URI, or POSTGRES_PASSWORD in .env and " + "`make scale-up SERVICE=pgvector`", "REQUIRE_PGVECTOR") schema = f"pytest_{uuid.uuid4().hex[:8]}" - adapter = PgvectorStore(uri, schema) + try: + adapter = PgvectorStore(uri, schema) + except VectorStoreError as exc: + _unavailable(str(exc), "REQUIRE_PGVECTOR") yield adapter adapter.drop_namespace() @@ -145,6 +154,10 @@ def test_the_callers_clearance_limits_what_is_returned(store): query = vectors[secret_idx].tolist() # identical to a secret chunk: it would rank first without the filter public = store.query("ops", query, len(ids), ["public"]) assert public and all(h.metadata["clearance"] == "public" for h in public) + # The store itself filters, rather than the guard cleaning up afterwards: asking for 3 of the 3 readable chunks + # returns all 3 even though the nearest chunks overall belong to higher clearances. + only_public = [i for i, m in enumerate(metas) if m["clearance"] == "public"] + assert len(store.query("ops", query, len(only_public), ["public"])) == len(only_public) assert ids[secret_idx] not in {h.id for h in public} internal = store.query("ops", query, len(ids), ["public", "internal"]) assert {h.metadata["clearance"] for h in internal} <= {"public", "internal"} and len(internal) > len(public) @@ -210,6 +223,22 @@ def test_a_reused_vector_ranks_like_a_fresh_one(store): assert [h.id for h in store.query("ops", [1.0] * DIM, 6, CLEARANCES)] == before +def test_results_match_chromas_own_output_not_only_numpy(store, chroma): + """Same corpus in Chroma and in the adapter under test: same chunks, same order, same similarity.""" + if isinstance(store, ChromaStore): + pytest.skip("this is Chroma") + publish(store, "ops", n=30, seed=4) + publish(chroma, "ops", n=30, seed=4) + rng = np.random.default_rng(11) + for allowed in (["public"], ["public", "internal"], CLEARANCES): + for _ in range(5): + query = (rng.normal(size=DIM) * 5.0).tolist() + theirs, reference = store.query("ops", query, 5, allowed), chroma.query("ops", query, 5, allowed) + assert [h.id for h in theirs] == [h.id for h in reference] + assert all(abs(a.similarity - b.similarity) < 1e-4 for a, b in zip(theirs, reference)) + assert [h.metadata["clearance"] for h in theirs] == [h.metadata["clearance"] for h in reference] + + # --- the guard belongs to the base class, not to the adapters ------------------------------------------------------------ class LeakyStore(VectorStore): @@ -289,3 +318,58 @@ def test_the_default_store_is_embedded_chroma_and_an_unknown_one_is_refused(monk monkeypatch.setattr(settings, "vector_store", "milvus") with pytest.raises(VectorStoreError, match="chroma, qdrant, pgvector"): get_vector_store() + + +# --- configuration errors say what to do ---------------------------------------------------------------------------------- + +def test_a_missing_qdrant_package_or_url_is_a_labelled_error(monkeypatch): + import sys + + monkeypatch.setattr(settings, "vector_store", "qdrant") + monkeypatch.setattr(settings, "qdrant_url", "http://localhost:6333") + monkeypatch.setitem(sys.modules, "qdrant_client", None) # as if it were not installed + monkeypatch.delitem(sys.modules, "rag_lab.storage.qdrant_store", raising=False) + vs._INSTANCES.clear() + with pytest.raises(VectorStoreError, match="qdrant-client"): + get_vector_store() + + +def test_qdrant_without_a_url_and_pgvector_without_a_uri_are_labelled_errors(): + pytest.importorskip("qdrant_client") + from rag_lab.storage.pgvector_store import PgvectorStore + from rag_lab.storage.qdrant_store import QdrantStore + + with pytest.raises(VectorStoreError, match="QDRANT_URL"): + QdrantStore("", "", "") + with pytest.raises(VectorStoreError, match="PGVECTOR_URI"): + PgvectorStore("", "public") + with pytest.raises(VectorStoreError, match="PGVECTOR_SCHEMA"): + PgvectorStore("postgresql://u:p@localhost/db", "x; DROP TABLE y") + + +def test_the_pgvector_pool_is_closed_at_exit_so_the_interpreter_finishes_cleanly(): + """Python 3.14 raises PythonFinalizationError when a pool's worker threads are still running at shutdown.""" + import subprocess + import sys + + import psycopg + from psycopg import sql + + from rag_lab.config import pgvector_uri_from_env + + uri = os.getenv("PGVECTOR_TEST_URI") or pgvector_uri_from_env() + if not uri: + _unavailable("no pgvector server", "REQUIRE_PGVECTOR") + schema = f"pytest_exit_{uuid.uuid4().hex[:6]}" + code = ( + "import sys\n" + "from rag_lab.storage.pgvector_store import PgvectorStore\n" + "PgvectorStore(sys.argv[1], sys.argv[2]).ensure('ws')\n" # then the interpreter exits: atexit must close the pool + ) + try: + done = subprocess.run([sys.executable, "-c", code, uri, schema], capture_output=True, text=True, + env=os.environ | {"PGVECTOR_URI": ""}) + finally: + with psycopg.connect(uri, autocommit=True) as conn: + conn.execute(sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(sql.Identifier(schema))) + assert done.returncode == 0 and "PythonFinalizationError" not in done.stderr, done.stderr[-400:] From 24efb2f03b46e7a6b70e65c6992f9bb7842dc9fa Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 20:37:16 +0200 Subject: [PATCH 3/8] feat(cache): the semantic cache can live in a shared PostgreSQL (CACHE_BACKEND=postgres) Off by default; SQLite is unchanged. With CACHE_DB_URI several UI processes read and write one cache. The rows are sealed exactly as in SQLite (AES-GCM, row identity as AAD), so the database operator sees no question, answer, citation or vector, and a row moved between clearances still refuses to open. The two statements that differ between the databases (schema and upsert) exist once per backend; everything else is shared with only the placeholders swapped, and a literal % is refused. A replaced row gets a new serial rowid, which the in-memory index uses to notice it. BYTEA comes back as bytes. A connection pool replaces a connection per call. L1 is off in shared mode because another process's prune or delete could not evict it. The existing cache suites (oracle comparison, tamper, review-flag, retention, proxies) now run against both backends through a cache_backend fixture; the dict-style proxy had its own INSERT OR REPLACE, which a new test now guards. --- rag_lab/retrieval/semantic_cache.py | 161 +++++++++++++++++++++++++--- tests/conftest.py | 22 ++++ tests/test_cache_index.py | 29 +++-- tests/test_semantic_cache.py | 27 +++-- tests/test_shared_cache.py | 118 ++++++++++++++++++++ 5 files changed, 323 insertions(+), 34 deletions(-) create mode 100644 tests/test_shared_cache.py diff --git a/rag_lab/retrieval/semantic_cache.py b/rag_lab/retrieval/semantic_cache.py index dcc0697..16cefe2 100644 --- a/rag_lab/retrieval/semantic_cache.py +++ b/rag_lab/retrieval/semantic_cache.py @@ -1,7 +1,10 @@ """Persistent, clearance-partitioned semantic and exact cache. Released answers are stored in a persistent SQLite database (CACHE_DB_PATH) with a thread-safe, -bounded LRU in-memory L1 cache. Cached answers are bound to workspace, clearance level, index +bounded LRU in-memory L1 cache. With CACHE_BACKEND=postgres the same table lives in a shared PostgreSQL database +(CACHE_DB_URI) so several UI processes see one cache; L1 is then off, because another process's prune or delete +could not evict it. The rows are sealed exactly the same way, so the database operator learns no question, answer, +citation or vector. All processes must share the same data key and the same index (docs/scaling.md). Cached answers are bound to workspace, clearance level, index manifest version, and query to prevent stale answers after re-indexing or unauthorized cross-clearance leakage. Query, answer, citations, and query vector are sealed at rest (AES-256-GCM, vault.py) with the row identity as additional authenticated data; rows that are not @@ -20,6 +23,7 @@ from __future__ import annotations import array +import atexit import collections import collections.abc import functools @@ -32,6 +36,7 @@ import sqlite3 import threading from datetime import UTC, datetime, timedelta +from urllib.parse import urlsplit from pathlib import Path from typing import Any, Generator @@ -171,9 +176,130 @@ def _open_row(row: tuple) -> dict[str, Any] | None: return None +def _shared() -> bool: + """True when the cache lives in PostgreSQL and is shared between processes.""" + return settings.cache_backend == "postgres" + + +# --- PostgreSQL backend --------------------------------------------------------------------------------------- +# Every statement in this module is a fixed template written with SQLite's `?` placeholders. The two statements that +# differ between the databases (the upsert and the schema) exist once per backend below; everything else is shared, +# with only the placeholders swapped. The `rowid` column is a serial: a replaced row gets a new one, which is how the +# in-memory index (cache_index.py) notices that a row was replaced. + +_PG_DDL = ( + """ + CREATE TABLE IF NOT EXISTS semantic_cache ( + key TEXT PRIMARY KEY, + workspace_id TEXT NOT NULL, + clearance TEXT NOT NULL, + index_version TEXT NOT NULL, + query TEXT NOT NULL, + generation TEXT NOT NULL, + citations_json TEXT NOT NULL, + query_vector BYTEA, + created_at TEXT NOT NULL, + reviewed INTEGER NOT NULL DEFAULT 1, + approved_by TEXT NOT NULL DEFAULT '', + rowid BIGSERIAL + ) + """, + "CREATE INDEX IF NOT EXISTS idx_cache_ws_clr_ver ON semantic_cache (workspace_id, clearance, index_version)", + """ + CREATE INDEX IF NOT EXISTS idx_cache_partition_sync + ON semantic_cache (workspace_id, clearance, index_version, created_at DESC, reviewed, key) + WHERE query_vector IS NOT NULL + """, +) +_UPSERT_SQLITE = f"INSERT OR REPLACE INTO semantic_cache ({_COLUMNS}) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);" +_UPSERT_PG = ( + f"INSERT INTO semantic_cache ({_COLUMNS}) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) " + "ON CONFLICT (key) DO UPDATE SET workspace_id = EXCLUDED.workspace_id, clearance = EXCLUDED.clearance, " + "index_version = EXCLUDED.index_version, query = EXCLUDED.query, generation = EXCLUDED.generation, " + "citations_json = EXCLUDED.citations_json, query_vector = EXCLUDED.query_vector, reviewed = EXCLUDED.reviewed, " + "approved_by = EXCLUDED.approved_by, created_at = EXCLUDED.created_at, " + "rowid = nextval(pg_get_serial_sequence('semantic_cache', 'rowid'))" +) + + +class _PgCursor: + """The part of a SQLite cursor this module uses. PostgreSQL returns BYTEA as memoryview; callers want bytes.""" + + def __init__(self, cursor: Any) -> None: + self._cursor = cursor + + @property + def rowcount(self) -> int: + return int(self._cursor.rowcount) + + @staticmethod + def _row(row: tuple | None) -> tuple | None: + return None if row is None else tuple(bytes(v) if isinstance(v, memoryview) else v for v in row) + + def fetchone(self) -> tuple | None: + return self._row(self._cursor.fetchone()) + + def fetchall(self) -> list[tuple]: + return [self._row(r) for r in self._cursor.fetchall()] # type: ignore[misc] + + +class _PgConn: + """A PostgreSQL connection with the `execute` / `commit` surface of sqlite3.Connection that this module uses.""" + + def __init__(self, raw: Any) -> None: + self._raw = raw + + def execute(self, statement: str, params: Any = ()) -> _PgCursor: + if "%" in statement: + raise ValueError("a literal % in a cache statement would be read as a PostgreSQL placeholder") + return _PgCursor(self._raw.execute(statement.replace("?", "%s"), params)) + + def commit(self) -> None: + self._raw.commit() + + +_PG_POOLS: dict[str, Any] = {} +_PG_READY: set[str] = set() +_PG_LOCK = threading.Lock() + + +def _pg_pool() -> Any: + from psycopg_pool import ConnectionPool + + uri = settings.cache_db_uri + if not uri: + raise RuntimeError("CACHE_BACKEND=postgres needs CACHE_DB_URI (a PostgreSQL connection string).") + with _PG_LOCK: + pool = _PG_POOLS.get(uri) + if pool is None: + pool = ConnectionPool(conninfo=uri, min_size=1, max_size=max(2, settings.checkpoint_pool_size), open=True, + timeout=15, kwargs={"autocommit": False}) + atexit.register(pool.close) # its worker threads must be stopped before the interpreter finalises + _PG_POOLS[uri] = pool + if uri not in _PG_READY: + with pool.connection() as raw: + raw.execute("SELECT pg_advisory_xact_lock(727274)") # two processes starting together create it once + for statement in _PG_DDL: + raw.execute(statement) + _PG_READY.add(uri) + return pool + + +def _db_identity(db_path: str | None) -> str: + """What names this database in the in-memory index: the resolved file, or the server and database (no password).""" + if _shared(): + parts = urlsplit(settings.cache_db_uri) + return f"postgres://{parts.hostname}:{parts.port or 5432}{parts.path}" + return str(Path(db_path or settings.cache_db_path).resolve()) + + @contextmanager -def _open_db(db_path: str | None = None) -> Generator[sqlite3.Connection, None, None]: - """Context manager guaranteeing connection termination and WAL-mode configuration.""" +def _open_db(db_path: str | None = None) -> Generator[Any, None, None]: + """A connection to the cache database: SQLite (the file `db_path`, default CACHE_DB_PATH) or the shared PostgreSQL.""" + if _shared(): + with _pg_pool().connection() as raw: + yield _PgConn(raw) + return target = db_path or settings.cache_db_path Path(target).parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(target, timeout=10.0) @@ -229,6 +355,8 @@ def _retention_cutoff() -> str: # --- L1 (LRU) ---------------------------------------------------------------- def _l1_put(key: str, generation: str, citations: list[str], reviewed: bool = True) -> None: + if _shared(): + return # another process's prune or delete could not evict this copy with _L1_LOCK: if key in _L1_GENERATIONS: _L1_GENERATIONS.move_to_end(key) @@ -243,6 +371,8 @@ def _l1_put(key: str, generation: str, citations: list[str], reviewed: bool = Tr def _l1_get(key: str) -> tuple[bool, str, list[str]]: + if _shared(): + return False, "", [] with _L1_LOCK: if key in _L1_GENERATIONS: _L1_GENERATIONS.move_to_end(key) @@ -311,10 +441,7 @@ def store_cache( ) with _open_db(db_path) as conn: conn.execute( - f""" - INSERT OR REPLACE INTO semantic_cache ({_COLUMNS}) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?); - """, + _UPSERT_PG if _shared() else _UPSERT_SQLITE, (ckey, workspace_id, clearance, ver, sealed_query, sealed_gen, sealed_cites, sealed_vec, 1 if reviewed else 0, approved_by, now), ) @@ -390,7 +517,7 @@ def lookup_cache( # 2. Semantic vector similarity across the matching workspace and index version return _semantic_lookup( - conn, str(Path(db_path or settings.cache_db_path).resolve()), workspace_id, authorized, curr_version, + conn, _db_identity(db_path), workspace_id, authorized, curr_version, cleaned_query, cutoff, sim_threshold, include_unreviewed, ) except Exception as exc: @@ -399,7 +526,7 @@ def lookup_cache( def _load_for_index( - conn: sqlite3.Connection, workspace_id: str, clearance: str, version: str, keys: list[str] + conn: Any, workspace_id: str, clearance: str, version: str, keys: list[str] ) -> dict[str, tuple[np.ndarray, Any]]: """Open the rows of one partition that the index has not seen. Rows that do not open are left out.""" placeholders = ", ".join("?" for _ in keys) @@ -420,7 +547,7 @@ def _load_for_index( def _semantic_lookup( - conn: sqlite3.Connection, + conn: Any, db_key: str, workspace_id: str, authorized: list[str], @@ -546,12 +673,12 @@ def clear_cache(db_path: str | None = None) -> None: def cache_stats(db_path: str | None = None) -> dict[str, Any]: """Retrieve cache entry counts, workspace distribution, and on-disk size.""" - target = db_path or settings.cache_db_path - path = Path(target) - if not path.is_file(): + path = Path(db_path or settings.cache_db_path) + if not _shared() and not path.is_file(): return {"total_entries": 0, "workspaces": {}, "bytes": 0} total = 0 + size = 0 workspaces: dict[str, int] = {} try: with _open_db(db_path) as conn: @@ -561,17 +688,19 @@ def cache_stats(db_path: str | None = None) -> dict[str, Any]: "SELECT workspace_id, COUNT(*) FROM semantic_cache GROUP BY workspace_id;" ).fetchall() workspaces = {r[0]: r[1] for r in ws_rows if r[0]} + if _shared(): + size = int(conn.execute("SELECT pg_total_relation_size('semantic_cache');").fetchone()[0]) # type: ignore[index] except Exception as exc: logger.debug("Error fetching cache stats: %s", exc) return { "total_entries": total, "workspaces": workspaces, - "bytes": path.stat().st_size if path.exists() else 0, + "bytes": size if _shared() else (path.stat().st_size if path.exists() else 0), } -def _fetch_entry(conn: sqlite3.Connection, key: str) -> dict[str, Any] | None: +def _fetch_entry(conn: Any, key: str) -> dict[str, Any] | None: row = conn.execute(f"SELECT {_COLUMNS} FROM semantic_cache WHERE key = ?;", (key,)).fetchone() return _open_row(row) if row else None @@ -600,7 +729,7 @@ def __setitem__(self, key: str, value: Any) -> None: with _open_db() as conn: now = datetime.now(UTC).isoformat(timespec="seconds") conn.execute( - f"INSERT OR REPLACE INTO semantic_cache ({_COLUMNS}) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);", + _UPSERT_PG if _shared() else _UPSERT_SQLITE, (key, "", "public", "", sealed_query, sealed_gen, sealed_cites, None, 1, "", now), ) conn.commit() diff --git a/tests/conftest.py b/tests/conftest.py index d382c30..911daaa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -86,6 +86,28 @@ def _fresh_accelerator_state(): accel._state.update(calls=0, noted=False) +@pytest.fixture(params=["sqlite", "postgres"]) +def cache_backend(request, monkeypatch): + """Run a test against the SQLite cache and against the shared PostgreSQL cache (CACHE_BACKEND=postgres). + + The Postgres run uses the dedicated `_pytest` database of the `postgres_uri` fixture, so it skips (or, with + REQUIRE_POSTGRES=1, fails) exactly like the checkpointer tests do. The shared table is emptied around every test. + """ + from rag_lab.retrieval import cache_index + from rag_lab.retrieval import semantic_cache as sc + + if request.param == "postgres": + uri = request.getfixturevalue("postgres_uri") + monkeypatch.setattr(settings, "cache_backend", "postgres") + monkeypatch.setattr(settings, "cache_db_uri", uri) + sc.clear_cache() + cache_index.forget() + yield request.param + if request.param == "postgres": + sc.clear_cache() + cache_index.forget() + + @pytest.fixture(autouse=True) def _clean_cache_each_test(): from rag_lab.retrieval.semantic_cache import clear_cache diff --git a/tests/test_cache_index.py b/tests/test_cache_index.py index c778ac3..4c1ff32 100644 --- a/tests/test_cache_index.py +++ b/tests/test_cache_index.py @@ -1,6 +1,5 @@ """The semantic cache's vector index: same answers as the old linear scan, never across clearances, never stale.""" import random -import sqlite3 import time from datetime import UTC, datetime, timedelta @@ -47,7 +46,9 @@ def fake(monkeypatch): @pytest.fixture -def db(tmp_path): +def db(tmp_path, cache_backend): + """The cache database of a test: a file for SQLite; for PostgreSQL the path is ignored (one shared table that + conftest.cache_backend empties around every test). Every test that takes `db` runs against both backends.""" return str(tmp_path / "cache.db") @@ -155,8 +156,9 @@ def test_a_row_edited_or_moved_on_disk_leaves_the_matrix_and_never_matches(db, f sc.store_cache("ops", question("alpha"), "internal", "internal answer", db_path=db) fake.near(question("bravo"), question("alpha"), 0.0, seed=5) assert sc.lookup_cache("ops", question("bravo"), "internal", threshold=0.5, db_path=db)["hit"] - with sqlite3.connect(db) as conn: + with sc._open_db(db) as conn: conn.execute("UPDATE semantic_cache SET clearance = 'public'") # move it down: the seal binds the clearance + conn.commit() for clearance in ("public", "internal"): sc._l1_clear() assert sc.lookup_cache("ops", question("bravo"), clearance, threshold=0.0, db_path=db) == {"hit": False} @@ -167,10 +169,11 @@ def test_a_row_whose_vector_is_swapped_for_another_rows_does_not_open(db, fake): sc.store_cache("ops", question("charlie"), "public", "second", db_path=db) fake.near(question("bravo"), question("alpha"), 0.0, seed=6) assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.5, db_path=db)["generation"] == "first" - with sqlite3.connect(db) as conn: - a, c = [r[0] for r in conn.execute("SELECT query_vector FROM semantic_cache ORDER BY rowid")] - conn.execute("UPDATE semantic_cache SET query_vector = ? WHERE generation IS NOT NULL AND rowid = 1", (c,)) - conn.execute("UPDATE semantic_cache SET query_vector = ? WHERE rowid = 2", (a,)) + with sc._open_db(db) as conn: + (first_id, a), (second_id, c) = conn.execute("SELECT rowid, query_vector FROM semantic_cache ORDER BY rowid").fetchall() + conn.execute("UPDATE semantic_cache SET query_vector = ? WHERE rowid = ?", (c, first_id)) + conn.execute("UPDATE semantic_cache SET query_vector = ? WHERE rowid = ?", (a, second_id)) + conn.commit() sc._l1_clear() assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.0, db_path=db) == {"hit": False} @@ -185,8 +188,9 @@ def test_changes_by_other_writers_are_seen_on_the_next_lookup(db, fake): sc._l1_clear() assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.5, db_path=db)["generation"] == "v2" assert cache_index.stats()["entries"] == cache_index_before, "replaced, not duplicated" - with sqlite3.connect(db) as conn: + with sc._open_db(db) as conn: conn.execute("DELETE FROM semantic_cache") + conn.commit() sc._l1_clear() assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.0, db_path=db) == {"hit": False} sc.store_cache("ops", question("delta"), "public", "after a wipe", db_path=db) # rowids start again at 1 @@ -198,12 +202,14 @@ def test_a_review_flag_changed_in_place_is_honoured_without_reopening_the_row(db sc.store_cache("ops", question("alpha"), "public", "auto-approved", db_path=db, reviewed=False) fake.near(question("bravo"), question("alpha"), 0.0, seed=9) assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.5, db_path=db, include_unreviewed=False) == {"hit": False} - with sqlite3.connect(db) as conn: + with sc._open_db(db) as conn: conn.execute("UPDATE semantic_cache SET reviewed = 1") + conn.commit() got = sc.lookup_cache("ops", question("bravo"), "public", threshold=0.5, db_path=db, include_unreviewed=False) assert got["hit"] and got["reviewed"] is True - with sqlite3.connect(db) as conn: + with sc._open_db(db) as conn: conn.execute("UPDATE semantic_cache SET reviewed = 0") + conn.commit() assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.5, db_path=db, include_unreviewed=False) == {"hit": False} @@ -212,8 +218,9 @@ def test_expired_entries_leave_the_index(db, fake): fake.near(question("bravo"), question("alpha"), 0.0, seed=10) assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.5, db_path=db)["hit"] stale = (datetime.now(UTC) - timedelta(days=400)).isoformat(timespec="seconds") - with sqlite3.connect(db) as conn: + with sc._open_db(db) as conn: conn.execute("UPDATE semantic_cache SET created_at = ?", (stale,)) + conn.commit() assert sc.lookup_cache("ops", question("bravo"), "public", threshold=0.0, db_path=db) == {"hit": False} assert cache_index.stats()["entries"] == 0 diff --git a/tests/test_semantic_cache.py b/tests/test_semantic_cache.py index af64ff4..26b7972 100644 --- a/tests/test_semantic_cache.py +++ b/tests/test_semantic_cache.py @@ -19,6 +19,10 @@ ) +# Every test here runs against the SQLite cache and against the shared PostgreSQL cache (conftest.cache_backend). +pytestmark = pytest.mark.usefixtures("cache_backend") + + @pytest.fixture(autouse=True) def clean_cache(): clear_cache() @@ -246,7 +250,9 @@ def test_cache_stats_and_malformed_vector_handling(tmp_path: Path): assert st["bytes"] > 0 -def test_cache_encryption_at_rest_and_row_aad_binding(tmp_path: Path): +def test_cache_encryption_at_rest_and_row_aad_binding(tmp_path: Path, cache_backend): + from rag_lab.retrieval.semantic_cache import _open_db + db_file = str(tmp_path / "cache.db") secret_answer = "Project-Omega-Confidential-998811" @@ -259,17 +265,24 @@ def test_cache_encryption_at_rest_and_row_aad_binding(tmp_path: Path): db_path=db_file, ) - with closing(sqlite3.connect(db_file)) as conn, conn: - conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") - raw_bytes = Path(db_file).read_bytes() - assert secret_answer.encode("utf-8") not in raw_bytes - assert b"What is the Omega code?" not in raw_bytes + with _open_db(db_file) as conn: # every column, as the database holds it + stored = conn.execute("SELECT key, workspace_id, clearance, index_version, query, generation, citations_json, " + "query_vector FROM semantic_cache").fetchall() + held = b"".join(v if isinstance(v, bytes) else str(v).encode() for row in stored for v in row if v is not None) + assert secret_answer.encode("utf-8") not in held and b"What is the Omega code?" not in held + assert b"omega.md" not in held, "citations are sealed too" + if cache_backend == "sqlite": # and the file itself, WAL included + with closing(sqlite3.connect(db_file)) as conn, conn: + conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + raw_bytes = Path(db_file).read_bytes() + assert secret_answer.encode("utf-8") not in raw_bytes + assert b"What is the Omega code?" not in raw_bytes # Tamper with row clearance in SQLite directly (attempting to downgrade secret row to public): # AAD includes clearance, so decryption fails and the tampered row is refused! _L1_GENERATIONS.clear() _L1_CITATIONS.clear() - with closing(sqlite3.connect(db_file)) as conn, conn: + with _open_db(db_file) as conn: conn.execute("UPDATE semantic_cache SET clearance = 'public'") conn.commit() diff --git a/tests/test_shared_cache.py b/tests/test_shared_cache.py new file mode 100644 index 0000000..e57d6de --- /dev/null +++ b/tests/test_shared_cache.py @@ -0,0 +1,118 @@ +"""The shared PostgreSQL cache (CACHE_BACKEND=postgres): what differs from the SQLite cache. + +The behavioural suites (tests/test_semantic_cache.py, tests/test_cache_index.py) run against both backends through +the `cache_backend` fixture, including the oracle comparison and the tamper tests. This file covers what only exists +for the shared backend. +""" +import re +from pathlib import Path + +import pytest + +import rag_lab.retrieval.semantic_cache as sc +from rag_lab.config import settings +from rag_lab.retrieval import cache_index + +SOURCE = Path(sc.__file__).read_text(encoding="utf-8") + + +def test_sqlite_only_sql_stays_in_the_sqlite_constants(): + """A statement that only SQLite understands, outside its constant, fails silently on PostgreSQL (the cache's broad + exception handling swallows it): the dict-style proxy once did exactly that.""" + assert SOURCE.count("INSERT OR REPLACE") == 1 and "_UPSERT_SQLITE = f\"INSERT OR REPLACE" in SOURCE + assert "PRAGMA" not in SOURCE.split("def _open_db")[1].split("yield conn")[0].split("if _shared():")[1].split("target =")[0] + statements = re.findall(r'"((?:SELECT|UPDATE|DELETE|INSERT)[^"]*)"', SOURCE) + assert statements and not [s for s in statements if "%" in s], "a literal % would be read as a placeholder" + + +def test_the_postgres_pool_is_not_created_unless_the_backend_is_selected(): + assert settings.cache_backend == "sqlite" and not sc._PG_POOLS + + +def test_a_missing_uri_is_a_labelled_error(monkeypatch): + monkeypatch.setattr(settings, "cache_backend", "postgres") + monkeypatch.setattr(settings, "cache_db_uri", "") + with pytest.raises(RuntimeError, match="CACHE_DB_URI"): + with sc._open_db(): + pass + + +def test_the_index_is_named_by_server_and_database_never_by_password(monkeypatch): + monkeypatch.setattr(settings, "cache_backend", "postgres") + monkeypatch.setattr(settings, "cache_db_uri", "postgresql://user:s3cret@db.internal:6543/rag_cache") + identity = sc._db_identity(None) + assert identity == "postgres://db.internal:6543/rag_cache" and "s3cret" not in identity + + +def test_l1_is_off_in_shared_mode(monkeypatch): + sc._l1_clear() + monkeypatch.setattr(settings, "cache_backend", "postgres") + sc._l1_put("k", "plaintext answer", ["c"]) + assert sc._l1_get("k") == (False, "", []), "another process's prune or delete could not evict a local copy" + monkeypatch.setattr(settings, "cache_backend", "sqlite") + sc._l1_put("k", "plaintext answer", ["c"]) + assert sc._l1_get("k")[0] is True + sc._l1_clear() + + +# --- against a real PostgreSQL (skips like the checkpointer tests; REQUIRE_POSTGRES=1 makes it fail instead) ------------ + +@pytest.fixture +def shared(cache_backend, request): + if cache_backend != "postgres": + pytest.skip("this test is about the shared backend") + return cache_backend + + +def test_the_table_holds_only_sealed_content_and_a_replaced_row_gets_a_new_serial(shared): + sc.store_cache("ops", "What is the Omega code?", "secret", "Project-Omega-998811", ["omega.md"]) + with sc._open_db() as conn: + first = conn.execute("SELECT rowid, query, generation, citations_json, query_vector FROM semantic_cache").fetchall() + assert len(first) == 1 + held = b"".join(v if isinstance(v, bytes) else str(v).encode() for v in first[0][1:] if v is not None) + assert b"Omega" not in held and b"omega.md" not in held + assert isinstance(first[0][4], bytes), "BYTEA comes back as bytes, not memoryview" + + sc.store_cache("ops", "What is the Omega code?", "secret", "Project-Omega-v2", ["omega.md"]) + with sc._open_db() as conn: + second = conn.execute("SELECT rowid FROM semantic_cache").fetchall() + assert len(second) == 1 and second[0][0] > first[0][0], "the index detects a replacement by its new rowid" + + +def test_two_processes_share_one_cache_without_a_restart(shared): + """A second process is simulated by a second, independent index registry reading the same table.""" + sc.store_cache("ops", "Which rule covers expense limits?", "public", "answer from process A", ["a.md"]) + sc._l1_clear() + cache_index.forget() # process B starts with nothing in memory + hit = sc.lookup_cache("ops", "Which rule covers expense limits?", "public") + assert hit["hit"] and hit["generation"] == "answer from process A" and hit["hit_type"] == "exact" + sc.store_cache("ops", "Which rule covers travel limits?", "public", "answer written by B", ["b.md"]) + cache_index.forget() # and A, restarted, sees B's write + again = sc.lookup_cache("ops", "Which rule covers travel limits?", "public") + assert again["hit"] and again["generation"] == "answer written by B" + + +def test_clearing_one_workspace_is_visible_to_every_process(shared): + sc.store_cache("ops", "Where is the runbook?", "public", "in the wiki", ["w.md"]) + sc.store_cache("hr", "How many leave days?", "public", "twenty-five", ["h.md"]) + sc.clear_workspace_cache("ops") + assert sc.lookup_cache("ops", "Where is the runbook?", "public") == {"hit": False} + assert sc.lookup_cache("hr", "How many leave days?", "public")["hit"] + stats = sc.cache_stats() + assert stats["total_entries"] == 1 and stats["workspaces"] == {"hr": 1} and stats["bytes"] > 0 + + +def test_concurrent_writers_do_not_lose_rows(shared): + import threading + + def writer(n): + for i in range(5): + sc.store_cache("ops", f"What is concurrent question {chr(97 + n)}{chr(97 + i)}?", "public", f"answer {n}{i}", ["c.md"], + compute_vector=False) + + threads = [threading.Thread(target=writer, args=(n,)) for n in range(6)] + for t in threads: + t.start() + for t in threads: + t.join() + assert sc.cache_stats()["total_entries"] == 30 From beb99799a2c9801e3eb585f5d8f3e6285426c166 Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 20:41:01 +0200 Subject: [PATCH 4/8] feat(llm): vLLM and other OpenAI-compatible servers are first-class: key, model, thinking, status LLM_PROVIDER=openai already existed but could not serve vLLM properly. Now: LOCAL_LLM_API_KEY (a secret, out of repr) instead of a hard-coded 'local'; LOCAL_LLM_MODEL instead of borrowing OLLAMA_MODEL; Qwen3's thinking, which is on by default in vLLM, is turned off with chat_template_kwargs (LOCAL_LLM_TEMPLATE_KWARGS, empty for servers that reject unknown fields), and any block left in the text, or a cut-off one, is removed so it never reaches the graders; reasoning fields are ignored and a null answer is empty rather than the string 'None'. make status probes GET /models and says unreachable / key refused / model not served. Tested against a stub server speaking the vLLM API over a Unix domain socket (runs in both the online and the offline suite): request shape, bearer auth, thinking in every form, error paths, a real generate_node answer, and the labelled extractive fallback when the server fails. --- rag_lab/pipeline/llm_provider.py | 86 ++++++++- rag_lab/storage/workspace_manager.py | 16 +- tests/test_openai_server.py | 258 +++++++++++++++++++++++++++ tests/test_workspace_manager.py | 31 ++++ 4 files changed, 385 insertions(+), 6 deletions(-) create mode 100644 tests/test_openai_server.py diff --git a/rag_lab/pipeline/llm_provider.py b/rag_lab/pipeline/llm_provider.py index 2e0b9f7..a041d6a 100644 --- a/rag_lab/pipeline/llm_provider.py +++ b/rag_lab/pipeline/llm_provider.py @@ -7,8 +7,10 @@ """ from __future__ import annotations +import json import logging import re +from dataclasses import dataclass from typing import Any import httpx @@ -111,6 +113,30 @@ def _llm_type(self) -> str: return "deterministic-context-synthesizer" +# A reasoning model's thinking, if a server leaves it in the text: it must never reach the graders or the reader. +_THINK_BLOCK = re.compile(r".*?", re.DOTALL | re.IGNORECASE) +_THINK_OPEN = re.compile(r".*", re.DOTALL | re.IGNORECASE) + + +def strip_thinking(text: str) -> str: + """Remove ... blocks, and a block that never closed (a cut-off answer is not an answer).""" + return _THINK_OPEN.sub("", _THINK_BLOCK.sub("", text)).strip() + + +def _template_kwargs() -> dict[str, Any]: + """LOCAL_LLM_TEMPLATE_KWARGS as a dict. Empty means send none, for servers that reject unknown request fields.""" + raw = settings.local_llm_template_kwargs.strip() + if not raw: + return {} + try: + parsed = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"LOCAL_LLM_TEMPLATE_KWARGS must be a JSON object or empty ({exc}).") from exc + if not isinstance(parsed, dict): + raise ValueError("LOCAL_LLM_TEMPLATE_KWARGS must be a JSON object or empty.") + return parsed + + def _openai_role(message: BaseMessage) -> str: if isinstance(message, SystemMessage): return "system" @@ -127,10 +153,13 @@ class OpenAICompatibleChat(BaseChatModel): base_url: str model_name: str = "qwen3:8b" - api_key: str = "local" + # A secret: kept out of repr() so it cannot reach a log line or a traceback by accident. + api_key: str = Field(default="local", repr=False) # Read when the client is created, so a settings override applies. timeout_sec: float = Field(default_factory=lambda: settings.ollama_timeout_sec) transport: Any = None + # Merged into the request as chat_template_kwargs; the default turns Qwen3's thinking off (LOCAL_LLM_TEMPLATE_KWARGS). + template_kwargs: dict[str, Any] = Field(default_factory=_template_kwargs) @property def _llm_type(self) -> str: @@ -143,11 +172,13 @@ def _generate( run_manager: Any = None, **kwargs: Any, ) -> ChatResult: - payload = { + payload: dict[str, Any] = { "model": self.model_name, "temperature": 0.0, "messages": [{"role": _openai_role(m), "content": str(m.content)} for m in messages], } + if self.template_kwargs: + payload["chat_template_kwargs"] = self.template_kwargs with httpx.Client(timeout=self.timeout_sec, transport=self.transport) as client: response = client.post( f"{self.base_url.rstrip('/')}/chat/completions", @@ -155,8 +186,53 @@ def _generate( headers={"Authorization": f"Bearer {self.api_key}"}, ) response.raise_for_status() - content = response.json()["choices"][0]["message"]["content"] - return ChatResult(generations=[ChatGeneration(message=AIMessage(content=str(content)))]) + choices = response.json().get("choices") or [] + if not choices: + raise ValueError("The chat server returned no choices.") + # Only `content` is the answer. A server that separates its reasoning (reasoning / reasoning_content) is ignored + # on purpose, and any block left in the text is removed. + content = choices[0].get("message", {}).get("content") or "" + return ChatResult(generations=[ChatGeneration(message=AIMessage(content=strip_thinking(str(content))))]) + + +@dataclass(frozen=True) +class OpenAIProbe: + models: set[str] | None + error: str = "" + + +def probe_openai(base_url: str | None = None, api_key: str | None = None, timeout_sec: float = 2.0, + transport: Any = None) -> OpenAIProbe: + """Ask an OpenAI-compatible server which models it serves (GET /models). Never raises; `error` says why not.""" + url = f"{(base_url or settings.local_llm_url).rstrip('/')}/models" + key = settings.local_llm_api_key if api_key is None else api_key + try: + with httpx.Client(timeout=timeout_sec, transport=transport) as client: + response = client.get(url, headers={"Authorization": f"Bearer {key}"}) + except httpx.HTTPError as exc: + return OpenAIProbe(None, f"unreachable ({type(exc).__name__})") + if response.status_code in (401, 403): + return OpenAIProbe(None, f"HTTP {response.status_code}: the server refused LOCAL_LLM_API_KEY") + if response.status_code != 200: + return OpenAIProbe(None, f"HTTP {response.status_code}") + try: + return OpenAIProbe({str(m["id"]) for m in response.json().get("data", [])}) + except (ValueError, KeyError, TypeError, AttributeError): + return OpenAIProbe(None, "an answer that is not an OpenAI model list") + + +def openai_model_name() -> str: + """The model id sent to an OpenAI-compatible server: LOCAL_LLM_MODEL, else the shared chat model setting.""" + return settings.local_llm_model or settings.chat_model + + +def _openai_chat(transport: Any = None) -> BaseChatModel: + return OpenAICompatibleChat( + base_url=settings.local_llm_url, + model_name=openai_model_name(), + api_key=settings.local_llm_api_key, + transport=transport, + ) def is_ollama_alive(base_url: str | None = None, timeout_sec: float = 0.5) -> bool: @@ -222,7 +298,7 @@ def get_chat_llm(provider: str | None = None) -> BaseChatModel: active = (provider or settings.llm_provider).lower().strip() if active == "openai": - return OpenAICompatibleChat(base_url=settings.local_llm_url, model_name=settings.chat_model) + return _openai_chat() if active == "deterministic": return DeterministicContextLLM() diff --git a/rag_lab/storage/workspace_manager.py b/rag_lab/storage/workspace_manager.py index 9e45d86..2bdfb30 100644 --- a/rag_lab/storage/workspace_manager.py +++ b/rag_lab/storage/workspace_manager.py @@ -382,7 +382,21 @@ def print_status() -> int: ok = False pulled = ollama_models(settings.ollama_base_url) chat_provider = settings.llm_provider.lower().strip() - if pulled is None: + if chat_provider == "openai": + from rag_lab.pipeline.llm_provider import openai_model_name, probe_openai + + probe = probe_openai() + wanted = openai_model_name() + if probe.models is None: + print(f"[!!] LLM backend : OpenAI-compatible server at {settings.local_llm_url}: {probe.error}") + ok = False + elif wanted in probe.models: + print(f"[OK] LLM backend : OpenAI-compatible server reachable ({wanted} at {settings.local_llm_url})") + else: + print(f"[!!] LLM backend : {settings.local_llm_url} is up but does not serve '{wanted}' " + f"(serves: {', '.join(sorted(probe.models)) or 'nothing'}); set LOCAL_LLM_MODEL") + ok = False + elif pulled is None: if chat_provider == "ollama": print(f"[!!] Ollama : offline at {settings.ollama_base_url}, but the configuration requires it") ok = False diff --git a/tests/test_openai_server.py b/tests/test_openai_server.py new file mode 100644 index 0000000..e5cb82e --- /dev/null +++ b/tests/test_openai_server.py @@ -0,0 +1,258 @@ +"""LLM_PROVIDER=openai against a stub of a vLLM-style server, over a Unix domain socket. + +A real HTTP server speaks the parts of the OpenAI API this provider uses (GET /v1/models, POST /v1/chat/completions, +bearer-token authentication, a reasoning model's habits). The socket is a Unix domain one, which the offline test run +allows, so these tests run in both suites without opening a TCP port. +""" +import http.server +import json +import shutil +import socketserver +import tempfile +import threading +from pathlib import Path + +import httpx +import pytest +from langchain_core.messages import HumanMessage, SystemMessage + +from rag_lab.config import settings +from rag_lab.pipeline import llm_provider +from rag_lab.pipeline.generation import generate_node +from rag_lab.pipeline.llm_provider import ( + OpenAICompatibleChat, + get_chat_llm, + openai_model_name, + probe_openai, + strip_thinking, +) + +MODEL = "Qwen/Qwen3-8B" +KEY = "sk-stub-secret-4711" + + +class StubServer: + """What vLLM does that matters here: bearer auth, a model list, and chat completions.""" + + def __init__(self, directory: str, api_key: str = KEY, models=(MODEL,)) -> None: + self.path = str(Path(directory) / "s.sock") + self.api_key = api_key + self.models = list(models) + self.requests: list[dict] = [] + self.reply = lambda payload: (200, {"choices": [{"message": {"content": "Five minutes."}}]}) + stub = self + + class Handler(http.server.BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *args): # a Unix-socket peer has no address to log + pass + + def _authorised(self) -> bool: + return not stub.api_key or self.headers.get("Authorization") == f"Bearer {stub.api_key}" + + def _send(self, status: int, body: dict) -> None: + data = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def do_GET(self): + stub.requests.append({"method": "GET", "path": self.path, "auth": self.headers.get("Authorization")}) + if not self._authorised(): + return self._send(401, {"error": "unauthorized"}) + self._send(200, {"object": "list", "data": [{"id": m, "object": "model"} for m in stub.models]}) + + def do_POST(self): + payload = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + stub.requests.append({"method": "POST", "path": self.path, "auth": self.headers.get("Authorization"), + "body": payload}) + if not self._authorised(): + return self._send(401, {"error": "unauthorized"}) + status, body = stub.reply(payload) + self._send(status, body) + + class Server(socketserver.ThreadingUnixStreamServer): + daemon_threads = True + + self._server = Server(self.path, Handler) + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + + @property + def transport(self) -> httpx.HTTPTransport: + return httpx.HTTPTransport(uds=self.path) + + def chat(self, **kw) -> OpenAICompatibleChat: + options = {"base_url": "http://vllm.test/v1", "model_name": MODEL, "api_key": KEY, "transport": self.transport} | kw + return OpenAICompatibleChat(**options) + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + + +@pytest.fixture +def server(): + directory = tempfile.mkdtemp(prefix="rs-", dir="/tmp") # a Unix socket path must be short + stub = StubServer(directory) + yield stub + stub.close() + shutil.rmtree(directory, ignore_errors=True) + + +@pytest.fixture +def configured(server, monkeypatch): + """Settings as an operator would set them for LLM_PROVIDER=openai, with the transport pointed at the stub.""" + monkeypatch.setattr(settings, "llm_provider", "openai") + monkeypatch.setattr(settings, "local_llm_url", "http://vllm.test/v1") + monkeypatch.setattr(settings, "local_llm_model", MODEL) + monkeypatch.setattr(settings, "local_llm_api_key", KEY) + original = llm_provider._openai_chat + monkeypatch.setattr(llm_provider, "_openai_chat", lambda transport=None: original(transport=server.transport)) + return server + + +# --- the request ----------------------------------------------------------------------------------------------------- + +def test_the_request_is_a_chat_completion_with_the_key_the_model_and_thinking_off(server): + reply = server.chat().invoke([SystemMessage(content="Be brief."), HumanMessage(content="Rollback SLA?")]) + assert reply.content == "Five minutes." + request = server.requests[-1] + assert request["path"] == "/v1/chat/completions" and request["auth"] == f"Bearer {KEY}" + body = request["body"] + assert body["model"] == MODEL and body["temperature"] == 0.0 + assert [m["role"] for m in body["messages"]] == ["system", "user"] + assert body["chat_template_kwargs"] == {"enable_thinking": False}, "Qwen3 thinks by default on vLLM" + + +def test_settings_drive_the_provider_the_factory_builds(configured): + llm = get_chat_llm() + assert isinstance(llm, OpenAICompatibleChat) + assert llm.invoke([HumanMessage(content="hi")]).content == "Five minutes." + assert configured.requests[-1]["body"]["model"] == MODEL and configured.requests[-1]["auth"] == f"Bearer {KEY}" + + +def test_the_model_id_is_its_own_setting_falling_back_to_the_shared_one(monkeypatch): + monkeypatch.setattr(settings, "local_llm_model", "") + monkeypatch.setattr(settings, "chat_model", "qwen3:8b") + assert openai_model_name() == "qwen3:8b" + monkeypatch.setattr(settings, "local_llm_model", MODEL) + assert openai_model_name() == MODEL + + +def test_template_kwargs_are_configurable_and_can_be_left_out(server, monkeypatch): + monkeypatch.setattr(settings, "local_llm_template_kwargs", "") + server.chat(template_kwargs=llm_provider._template_kwargs()).invoke([HumanMessage(content="hi")]) + assert "chat_template_kwargs" not in server.requests[-1]["body"], "for a server that rejects unknown fields" + + monkeypatch.setattr(settings, "local_llm_template_kwargs", '{"thinking": false}') + server.chat(template_kwargs=llm_provider._template_kwargs()).invoke([HumanMessage(content="hi")]) + assert server.requests[-1]["body"]["chat_template_kwargs"] == {"thinking": False} + + +@pytest.mark.parametrize("bad", ["not json", "[1, 2]", '"text"']) +def test_bad_template_kwargs_are_a_labelled_error(monkeypatch, bad): + monkeypatch.setattr(settings, "local_llm_template_kwargs", bad) + with pytest.raises(ValueError, match="LOCAL_LLM_TEMPLATE_KWARGS"): + llm_provider._template_kwargs() + + +# --- the answer ------------------------------------------------------------------------------------------------------------ + +@pytest.mark.parametrize("returned,expected", [ + ("The user asks about rollback.\nFive minutes.", "Five minutes."), + ("x\nyFive minutes.", "Five minutes."), + ("cut off before it finished", ""), + ("Five minutes.", "Five minutes."), +]) +def test_thinking_left_in_the_text_never_reaches_the_reader(server, returned, expected): + server.reply = lambda payload: (200, {"choices": [{"message": {"content": returned}}]}) + assert server.chat().invoke([HumanMessage(content="hi")]).content == expected + + +def test_a_separate_reasoning_field_is_ignored_and_a_null_answer_is_empty_not_the_word_none(server): + server.reply = lambda payload: (200, {"choices": [{"message": {"content": "Five minutes.", "reasoning_content": "secret chain of thought", + "reasoning": "more"}}]}) + assert server.chat().invoke([HumanMessage(content="hi")]).content == "Five minutes." + server.reply = lambda payload: (200, {"choices": [{"message": {"content": None, "reasoning_content": "only thoughts"}}]}) + assert server.chat().invoke([HumanMessage(content="hi")]).content == "" + + +def test_strip_thinking_on_its_own(): + assert strip_thinking("a b c d e") == "a c e" + assert strip_thinking("") == "" + + +# --- failures ------------------------------------------------------------------------------------------------------------------ + +def test_the_wrong_key_is_an_http_error_that_does_not_echo_the_key(server): + with pytest.raises(httpx.HTTPStatusError) as caught: + server.chat(api_key="wrong-key").invoke([HumanMessage(content="hi")]) + assert caught.value.response.status_code == 401 and "wrong-key" not in str(caught.value) + + +def test_server_errors_and_empty_answers_are_errors(server): + server.reply = lambda payload: (500, {"error": "boom"}) + with pytest.raises(httpx.HTTPStatusError): + server.chat().invoke([HumanMessage(content="hi")]) + server.reply = lambda payload: (200, {"choices": []}) + with pytest.raises(ValueError, match="no choices"): + server.chat().invoke([HumanMessage(content="hi")]) + + +def test_an_unreachable_server_is_a_connection_error(tmp_path): + llm = OpenAICompatibleChat(base_url="http://vllm.test/v1", transport=httpx.HTTPTransport(uds="/tmp/no-such-socket")) + with pytest.raises(httpx.ConnectError): + llm.invoke([HumanMessage(content="hi")]) + + +def test_the_api_key_stays_out_of_repr(server): + assert KEY not in repr(server.chat()) and KEY not in str(server.chat()) + + +# --- an answer through the pipeline ---------------------------------------------------------------------------------------------- + +def test_generation_answers_through_the_server(configured): + configured.reply = lambda payload: (200, {"choices": [{"message": {"content": "checkingThe rollback SLA is under 5 minutes."}}]}) + state = generate_node({ + "documents": ["Rollback SLA is strictly under 5 minutes."], + "sanitized_question": "What is the rollback SLA?", + "degraded": [], + "telemetry": {}, + }) + assert state["generation"] == "The rollback SLA is under 5 minutes." and "" not in state["generation"] + assert not state["telemetry"].get("generation_fallback") + sent = configured.requests[-1]["body"]["messages"] + assert "Rollback SLA is strictly under 5 minutes." in json.dumps(sent), "the retrieved context reached the server" + + +def test_a_failing_server_gives_the_labelled_extractive_fallback(configured): + configured.reply = lambda payload: (503, {"error": "overloaded"}) + state = generate_node({ + "documents": ["Rollback SLA is strictly under 5 minutes."], + "sanitized_question": "What is the rollback SLA?", + "degraded": [], + "telemetry": {}, + }) + assert "5 minutes" in state["generation"] and state["telemetry"]["generation_fallback"] is True + assert any(note.startswith("generation_fallback: HTTPStatusError") for note in state["degraded"]) + + +# --- the status probe -------------------------------------------------------------------------------------------------------------- + +def test_probe_lists_the_served_models(server): + probe = probe_openai("http://vllm.test/v1", KEY, transport=server.transport) + assert probe.models == {MODEL} and probe.error == "" + assert server.requests[-1]["path"] == "/v1/models" + + +def test_probe_says_why_it_failed(server): + refused = probe_openai("http://vllm.test/v1", "wrong", transport=server.transport) + assert refused.models is None and "401" in refused.error and "LOCAL_LLM_API_KEY" in refused.error + gone = probe_openai("http://vllm.test/v1", KEY, transport=httpx.HTTPTransport(uds="/tmp/no-such-socket")) + assert gone.models is None and "unreachable" in gone.error + server.models = [] + assert probe_openai("http://vllm.test/v1", KEY, transport=server.transport).models == set() diff --git a/tests/test_workspace_manager.py b/tests/test_workspace_manager.py index dbaf1db..6910689 100644 --- a/tests/test_workspace_manager.py +++ b/tests/test_workspace_manager.py @@ -240,3 +240,34 @@ def test_status_shows_the_python_fallback_as_degraded_and_fails(rollback_index, assert code == 1 assert "[!!] Context sanitizer : DEGRADED, pure-Python fallback (No module named 'rag_accel')" in out assert "wheelhouse" in out and "make rust-build" in out + + +# --- make status: an OpenAI-compatible chat server (vLLM, LM Studio, llama.cpp) --------------------------------- + +def _openai_status(capsys, monkeypatch, probe): + from rag_lab.pipeline import llm_provider + + monkeypatch.setattr(settings, "local_llm_url", "http://vllm.test/v1") + monkeypatch.setattr(settings, "local_llm_model", "Qwen/Qwen3-8B") + monkeypatch.setattr(llm_provider, "probe_openai", lambda *a, **k: probe) + return _status(capsys, monkeypatch, llm="openai") + + +def test_status_reports_a_reachable_openai_server_that_serves_the_model(rollback_index, capsys, monkeypatch): + from rag_lab.pipeline.llm_provider import OpenAIProbe + + code, out = _openai_status(capsys, monkeypatch, OpenAIProbe({"Qwen/Qwen3-8B"})) + assert code == 0, out + assert "[OK] LLM backend : OpenAI-compatible server reachable (Qwen/Qwen3-8B at http://vllm.test/v1)" in out + assert "Ollama" not in out.split("LLM backend")[1].split("\n")[0] + + +def test_status_flags_an_openai_server_that_is_down_refuses_the_key_or_lacks_the_model(rollback_index, capsys, monkeypatch): + from rag_lab.pipeline.llm_provider import OpenAIProbe + + code, out = _openai_status(capsys, monkeypatch, OpenAIProbe(None, "unreachable (ConnectError)")) + assert code == 1 and "[!!] LLM backend" in out and "unreachable (ConnectError)" in out + code, out = _openai_status(capsys, monkeypatch, OpenAIProbe(None, "HTTP 401: the server refused LOCAL_LLM_API_KEY")) + assert code == 1 and "LOCAL_LLM_API_KEY" in out + code, out = _openai_status(capsys, monkeypatch, OpenAIProbe({"other/model"})) + assert code == 1 and "does not serve 'Qwen/Qwen3-8B'" in out and "set LOCAL_LLM_MODEL" in out From 1189c28b24fd9b80f31a4da1418d1bc77dccae9a Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 20:44:31 +0200 Subject: [PATCH 5/8] docs+ci: docs/scaling.md, the Qdrant and pgvector services in CI, and tests that pin the profiles docs/scaling.md covers each option (what it gives, what it costs, how to turn it on and off), the measurements (4,000 and 20,000 chunks: the stores differ by single milliseconds at the sizes this project targets), and what does not scale: the keyword index, uploads, the audit log, the policy and the keys stay local files, so several UI processes on one host work and several hosts need shared storage. CI starts the two services and fails (REQUIRE_QDRANT / REQUIRE_PGVECTOR) rather than skips if either is missing; the offline run skips only the Qdrant server tests, which need Python TCP sockets. A new test pins the compose shape (one profile each, loopback only, no secrets, vLLM asks for a GPU), the Makefile targets, and that every scaling variable is read by the code and documented. --- .env.example | 21 ++++++ .github/workflows/ci.yml | 27 +++++++ AGENTS.md | 2 + CLAUDE.md | 2 + README.md | 2 + docs/scaling.md | 137 +++++++++++++++++++++++++++++++++++ tests/test_scaling_config.py | 112 ++++++++++++++++++++++++++++ tests/test_vector_stores.py | 2 + 8 files changed, 305 insertions(+) create mode 100644 docs/scaling.md create mode 100644 tests/test_scaling_config.py diff --git a/.env.example b/.env.example index 4032c55..3be24e5 100644 --- a/.env.example +++ b/.env.example @@ -51,3 +51,24 @@ UPLOADS_DIR=./uploads # The native Rust accelerator (rag_accel): preferred (default) uses it and falls back, loudly, to pure Python when it # is missing; required refuses to start without it. 'make setup' installs it (wheelhouse/, then cargo). # RAG_ACCEL=preferred + +# --- Scaling options, all off by default (docs/scaling.md) -------------------------------------------------------- +# Where chunk vectors live: chroma (embedded, default), qdrant or pgvector. Switching needs a re-sync: +# make sync ADMIN= FORCE=1. Start a server with: make scale-up SERVICE=qdrant (or pgvector). +# VECTOR_STORE=chroma +# QDRANT_URL=http://localhost:6333 +# QDRANT_API_KEY= # empty = no authentication (loopback only); set one before exposing the port +# QDRANT_PREFIX= # prepended to collection names so installs can share a server +# PGVECTOR_URI= # default: built from POSTGRES_* for the compose profile (port 5433) +# PGVECTOR_SCHEMA=public + +# The semantic cache: sqlite (default) or postgres, shared by several UI processes. All of them need the same data key. +# CACHE_BACKEND=sqlite +# CACHE_DB_URI=postgresql://user:password@host:5432/dbname + +# A chat model on vLLM, LM Studio or llama.cpp: LLM_PROVIDER=openai +# LOCAL_LLM_URL=http://localhost:8000/v1 +# LOCAL_LLM_MODEL=Qwen/Qwen3-8B +# LOCAL_LLM_API_KEY= +# LOCAL_LLM_TEMPLATE_KWARGS={"enable_thinking": false} # empty = send none (a server that rejects unknown fields) +# VLLM_MODEL=Qwen/Qwen3-8B # for `make scale-up SERVICE=vllm` (needs an NVIDIA GPU); VLLM_API_KEY, VLLM_PORT diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d196527..d02ecdd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,6 +15,9 @@ jobs: # Tests run on throwaway SQLite stores; this server backs the PostgreSQL checkpointer tests. # Credentials of a throwaway service container that exists only for this run; real deployments use .env. POSTGRES_TEST_URI: postgresql://rag_user:rag_password@localhost:5432/rag_checkpoints + # The scaling options (docs/scaling.md) are tested against real servers; a missing one fails the run (REQUIRE_*). + QDRANT_TEST_URL: http://localhost:6333 + PGVECTOR_TEST_URI: postgresql://rag_user:rag_password@localhost:5433/rag_vectors services: postgres: @@ -31,6 +34,25 @@ jobs: --health-timeout 5s --health-retries 5 + qdrant: + image: qdrant/qdrant:v1.20.0 + ports: + - 6333:6333 + + pgvector: + image: pgvector/pgvector:pg16 + env: + POSTGRES_DB: rag_vectors + POSTGRES_USER: rag_user + POSTGRES_PASSWORD: rag_password + ports: + - 5433:5432 + options: >- + --health-cmd "pg_isready -U rag_user -d rag_vectors" + --health-interval 5s + --health-timeout 5s + --health-retries 5 + steps: - name: Check out repository uses: actions/checkout@v7 @@ -46,6 +68,7 @@ jobs: python -m pip install --upgrade pip pip install -r requirements.lock pip install maturin + pip install "qdrant-client>=1.9" # optional: the Qdrant adapter is tested here - name: Lint (ruff, rule set pinned in pyproject.toml) run: ruff check . @@ -75,6 +98,8 @@ jobs: # A missing PostgreSQL server or Rust accelerator fails the run instead of skipping its tests. REQUIRE_POSTGRES: "1" REQUIRE_RUST_ACCEL: "1" + REQUIRE_QDRANT: "1" + REQUIRE_PGVECTOR: "1" # The accelerator is the default here: a missing one stops the process instead of falling back. RAG_ACCEL: required run: pytest -v --cov=. --cov-report=term-missing:skip-covered @@ -84,5 +109,7 @@ jobs: RAG_OFFLINE: "1" REQUIRE_POSTGRES: "1" REQUIRE_RUST_ACCEL: "1" + REQUIRE_QDRANT: "1" + REQUIRE_PGVECTOR: "1" RAG_ACCEL: required run: pytest -v diff --git a/AGENTS.md b/AGENTS.md index 7c1be43..fa53bca 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -84,6 +84,7 @@ make clean # Clean caches and temporary test artifacts make factory-reset [CONFIRM=yes] make up | down | db-password # PostgreSQL (127.0.0.1, password in .env, created by `make up`); db-password applies it to an old volume make lan-tls | lan-tls-down | lan-tls-ca # optional Caddy HTTPS for the LAN (docs/lan-tls.md) +make scale-up | scale-down SERVICE= # optional scaling services (docs/scaling.md); `make up` starts none ``` --- @@ -125,6 +126,7 @@ Graph wiring is in [`rag_lab/pipeline/graph.py`](rag_lab/pipeline/graph.py); sta * **Keys & Second Factor**: [`keystore.py`](rag_lab/security/keystore.py) owns where the access and data keys live: a key file holds the raw key, a pointer to an OS-keychain entry (random account per key), or an envelope-wrapped key (interface + stub only). `KEY_STORE=auto|file|keyring`; existing files move only via `make migrate-keys`; a pointer whose entry is gone is an error, never a new key; keychain keys are cached per process. Tests pin `KEY_STORE=file` and an inert keyring backend: nothing may reach the real keychain. TOTP (`access.py`): self-enrolment in two steps, secret sealed with AAD `uid:name`, code spent under the policy lock for a newer time step only, recovery codes stored as keyed tags, enforced for as long as the enrolment exists. Lockouts: an owner's password and second factor are reset by an owner (`set_password`, `disable_totp`), a sole owner's by `recover_account` (CLI only, README table "When someone is locked out"). pyotp 2.10 has no `verify_and_get_timecode`: `_match_totp_step` is ours. * **Clearance-Partitioned Sealed Cache**: `SEMANTIC_CACHE` and `SEMANTIC_CACHE_CITATIONS` ([`semantic_cache.py`](rag_lab/retrieval/semantic_cache.py)) provide an LRU in-memory + AES-256-GCM sealed SQLite-persistent cache keyed on `(workspace_id, clearance, index_fingerprint, original_question)` with cosine-similarity semantic retrieval, numeric/negation signature verification, and `reviewed` status tracking so users without `auto_approve` only consume human-reviewed cache entries. Semantic search is [`cache_index.py`](rag_lab/retrieval/cache_index.py): one float32 matrix per (db, workspace, clearance, index version) partition; a lookup searches only the caller's authorized clearances, re-syncs each partition with the table by `(key, rowid, created_at, reviewed)` first, and the winning row is re-read, opened (seal + AAD), re-checked and its cosine recomputed before a hit. The index only proposes candidates; never make it the authority. * **Fail-Closed Document Clearance**: Unrecognized explicit `clearance` tags in YAML frontmatter (case-insensitive) or `.clearance` folder files are rejected during ingestion ([`ingest.py`](rag_lab/ingestion/ingest.py)). Both dense ([`retrieval.py`](rag_lab/pipeline/retrieval.py)) and sparse ([`sparse_search.py`](rag_lab/retrieval/sparse_search.py)) retrieval enforce `allowed_clearances` at query time and via fail-closed post-query guards that reject chunks missing a valid clearance attribute. +* **Scaling options (off by default)**: all vector access goes through [`VectorStore`](rag_lab/storage/vector_store.py) (`VECTOR_STORE=chroma|qdrant|pgvector`); `query()` applies `guard_hits()` after the adapter's server-side filter, so the clearance guard lives in the base class and a new adapter cannot skip it. Similarity must stay cosine on Chroma's `1 - distance` scale (the gates are calibrated on it: tests compare values, not just order). Qdrant/pgvector fail closed (no unfiltered fallback) and publish atomically. `CACHE_BACKEND=postgres` puts the sealed semantic cache in a shared database (L1 off; same data key everywhere). `LLM_PROVIDER=openai` (`LOCAL_LLM_*`) serves vLLM/LM Studio/llama.cpp: thinking off via `chat_template_kwargs`, `` stripped. Tests pin all of it off (conftest, eval redirect) so a `.env` can never make a test reach a real service; adapters are tested in their own namespace (`QDRANT_PREFIX`, `PGVECTOR_SCHEMA`). * **Ingest reuse**: a re-sync embeds only chunks whose exact text is new and takes the stored vector of the rest, but only from the same workspace's live collection with the same model fingerprint (collection and manifest) and the same `passage_prefix` recorded in the manifest; `FORCE=1`, an unreadable index or an older manifest embeds everything. Plain-text files (`PLAIN_SUFFIXES`) are decoded in-process; every other format stays in the sandboxed parser child. * **Workspaces**: Sanitized slugs (`ws_`). Each workspace deterministically merges `knowledge_base/` and `uploads/` (with `knowledge_base` taking precedence on filename collisions) and performs an atomic Chroma collection swap (`__build` -> live) during `FORCE=1` re-syncs. Deleting a workspace requires admin privileges, escapes SQL `LIKE` wildcards, and prunes its grants from `access_policy.json`. * **Human-in-the-Loop & Safety Escalation**: Answers release automatically when the user holds `auto_approve` **and** the generation is grounded, relevant, non-abstaining, and free of blocking degradation (`security_warning`, `dense_retrieval_unavailable`, `sparse_retrieval_unavailable`). Otherwise, execution pauses at `interrupt()` for a reviewer who holds `review` permission, is not the requester (`requester != approver`), and whose clearance dominates (`clearance_covers`) the thread's `principal_clearance`. diff --git a/CLAUDE.md b/CLAUDE.md index 7c1be43..fa53bca 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -84,6 +84,7 @@ make clean # Clean caches and temporary test artifacts make factory-reset [CONFIRM=yes] make up | down | db-password # PostgreSQL (127.0.0.1, password in .env, created by `make up`); db-password applies it to an old volume make lan-tls | lan-tls-down | lan-tls-ca # optional Caddy HTTPS for the LAN (docs/lan-tls.md) +make scale-up | scale-down SERVICE= # optional scaling services (docs/scaling.md); `make up` starts none ``` --- @@ -125,6 +126,7 @@ Graph wiring is in [`rag_lab/pipeline/graph.py`](rag_lab/pipeline/graph.py); sta * **Keys & Second Factor**: [`keystore.py`](rag_lab/security/keystore.py) owns where the access and data keys live: a key file holds the raw key, a pointer to an OS-keychain entry (random account per key), or an envelope-wrapped key (interface + stub only). `KEY_STORE=auto|file|keyring`; existing files move only via `make migrate-keys`; a pointer whose entry is gone is an error, never a new key; keychain keys are cached per process. Tests pin `KEY_STORE=file` and an inert keyring backend: nothing may reach the real keychain. TOTP (`access.py`): self-enrolment in two steps, secret sealed with AAD `uid:name`, code spent under the policy lock for a newer time step only, recovery codes stored as keyed tags, enforced for as long as the enrolment exists. Lockouts: an owner's password and second factor are reset by an owner (`set_password`, `disable_totp`), a sole owner's by `recover_account` (CLI only, README table "When someone is locked out"). pyotp 2.10 has no `verify_and_get_timecode`: `_match_totp_step` is ours. * **Clearance-Partitioned Sealed Cache**: `SEMANTIC_CACHE` and `SEMANTIC_CACHE_CITATIONS` ([`semantic_cache.py`](rag_lab/retrieval/semantic_cache.py)) provide an LRU in-memory + AES-256-GCM sealed SQLite-persistent cache keyed on `(workspace_id, clearance, index_fingerprint, original_question)` with cosine-similarity semantic retrieval, numeric/negation signature verification, and `reviewed` status tracking so users without `auto_approve` only consume human-reviewed cache entries. Semantic search is [`cache_index.py`](rag_lab/retrieval/cache_index.py): one float32 matrix per (db, workspace, clearance, index version) partition; a lookup searches only the caller's authorized clearances, re-syncs each partition with the table by `(key, rowid, created_at, reviewed)` first, and the winning row is re-read, opened (seal + AAD), re-checked and its cosine recomputed before a hit. The index only proposes candidates; never make it the authority. * **Fail-Closed Document Clearance**: Unrecognized explicit `clearance` tags in YAML frontmatter (case-insensitive) or `.clearance` folder files are rejected during ingestion ([`ingest.py`](rag_lab/ingestion/ingest.py)). Both dense ([`retrieval.py`](rag_lab/pipeline/retrieval.py)) and sparse ([`sparse_search.py`](rag_lab/retrieval/sparse_search.py)) retrieval enforce `allowed_clearances` at query time and via fail-closed post-query guards that reject chunks missing a valid clearance attribute. +* **Scaling options (off by default)**: all vector access goes through [`VectorStore`](rag_lab/storage/vector_store.py) (`VECTOR_STORE=chroma|qdrant|pgvector`); `query()` applies `guard_hits()` after the adapter's server-side filter, so the clearance guard lives in the base class and a new adapter cannot skip it. Similarity must stay cosine on Chroma's `1 - distance` scale (the gates are calibrated on it: tests compare values, not just order). Qdrant/pgvector fail closed (no unfiltered fallback) and publish atomically. `CACHE_BACKEND=postgres` puts the sealed semantic cache in a shared database (L1 off; same data key everywhere). `LLM_PROVIDER=openai` (`LOCAL_LLM_*`) serves vLLM/LM Studio/llama.cpp: thinking off via `chat_template_kwargs`, `` stripped. Tests pin all of it off (conftest, eval redirect) so a `.env` can never make a test reach a real service; adapters are tested in their own namespace (`QDRANT_PREFIX`, `PGVECTOR_SCHEMA`). * **Ingest reuse**: a re-sync embeds only chunks whose exact text is new and takes the stored vector of the rest, but only from the same workspace's live collection with the same model fingerprint (collection and manifest) and the same `passage_prefix` recorded in the manifest; `FORCE=1`, an unreadable index or an older manifest embeds everything. Plain-text files (`PLAIN_SUFFIXES`) are decoded in-process; every other format stays in the sandboxed parser child. * **Workspaces**: Sanitized slugs (`ws_`). Each workspace deterministically merges `knowledge_base/` and `uploads/` (with `knowledge_base` taking precedence on filename collisions) and performs an atomic Chroma collection swap (`__build` -> live) during `FORCE=1` re-syncs. Deleting a workspace requires admin privileges, escapes SQL `LIKE` wildcards, and prunes its grants from `access_policy.json`. * **Human-in-the-Loop & Safety Escalation**: Answers release automatically when the user holds `auto_approve` **and** the generation is grounded, relevant, non-abstaining, and free of blocking degradation (`security_warning`, `dense_retrieval_unavailable`, `sparse_retrieval_unavailable`). Otherwise, execution pauses at `interrupt()` for a reviewer who holds `review` permission, is not the requester (`requester != approver`), and whose clearance dominates (`clearance_covers`) the thread's `principal_clearance`. diff --git a/README.md b/README.md index 79c71f0..bcf9b2f 100644 --- a/README.md +++ b/README.md @@ -374,6 +374,8 @@ The repository includes a `docker-compose.yml` to run PostgreSQL (`postgres:16-a **Credentials.** There is no password in the repository. `make up` creates `.env` (mode 600, git-ignored) with a random `POSTGRES_PASSWORD` the first time; `docker-compose.yml` reads it through `${POSTGRES_PASSWORD}` and refuses to start without one. The server listens on `127.0.0.1` only. To use it as the checkpoint store, uncomment the `DB_URI` line in `.env` (it expands `${POSTGRES_USER}` and `${POSTGRES_PASSWORD}`). The PostgreSQL tests use `POSTGRES_TEST_URI`, or the credentials in `.env`. +**Scaling options (all off by default).** Qdrant or pgvector as the vector store, vLLM or any OpenAI-compatible server as the chat model, a PostgreSQL-backed semantic cache shared by several UI processes, each a compose profile started with `make scale-up SERVICE=` that `make up` never starts: see [docs/scaling.md](docs/scaling.md), including the measurements and what does **not** scale (the keyword index, the audit log and the policy stay local files). + **HTTPS on a LAN.** `make ui` serves plain HTTP. For colleagues on a network there is an optional Caddy proxy with a local CA (`make lan-tls`, a compose profile that `make up` never starts): see [docs/lan-tls.md](docs/lan-tls.md). **Upgrading from an earlier version.** The old compose file used a fixed password, and PostgreSQL applies `POSTGRES_PASSWORD` only when a volume is first created. After `make up` has created `.env`, run `make db-password` once: it sets the existing role's password to the one in `.env` over the container's local socket (the data is untouched). Starting fresh instead (`docker compose down -v`) deletes the volume and its checkpoints. diff --git a/docs/scaling.md b/docs/scaling.md new file mode 100644 index 0000000..04be85a --- /dev/null +++ b/docs/scaling.md @@ -0,0 +1,137 @@ +# Scaling options (all off by default) + +Nothing here is needed on one laptop: `make setup && make models && make ui` uses embedded Chroma, a SQLite semantic +cache, and Ollama or the in-process synthesizer. Each option below is opt-in, starts nothing on its own, and can be +turned off again by deleting one line of `.env`. Read *What does not scale* before you plan a multi-host deployment. + +| Option | Turn it on | What it gives you | Cost | +|---|---|---|---| +| Vector store: **Qdrant** | `make scale-up SERVICE=qdrant`, `VECTOR_STORE=qdrant` | One vector index shared by several processes or hosts; 3x faster queries and 4x faster publishes than embedded Chroma at 20,000 chunks | A server (292 MB image), `pip install qdrant-client`, a re-sync | +| Vector store: **pgvector** | `make scale-up SERVICE=pgvector`, `VECTOR_STORE=pgvector` | The same, inside PostgreSQL you may already run and back up | A server (640 MB image), a re-sync; exact scan, no ANN index | +| Chat model: **vLLM** / any OpenAI-compatible server | `LLM_PROVIDER=openai`, `LOCAL_LLM_URL=...` (`make scale-up SERVICE=vllm` on a GPU host) | Throughput for many concurrent users; a bigger model than one laptop runs | An NVIDIA GPU for vLLM; the model is not hash-pinned | +| **Shared semantic cache** | `CACHE_BACKEND=postgres`, `CACHE_DB_URI=...` | Several UI processes see one cache, sealed as before | A PostgreSQL database; a few conditions on keys and indexes (below) | +| **LAN TLS** | `make lan-tls` | HTTPS in front of the UI | See [lan-tls.md](lan-tls.md) | + +`make scale-up SERVICE=` starts one compose profile; `make scale-down SERVICE=...` stops it and keeps +its data. `make up` never starts any of them. Every service binds `127.0.0.1`. + +## Vector store + +Chroma stays the default. `VECTOR_STORE=qdrant` or `pgvector` swaps where chunk vectors live, behind one interface +(`rag_lab/storage/vector_store.py`), and **nothing else in the pipeline changes**. Switching does not migrate: the old +index is left where it was, and `make sync ADMIN= FORCE=1` builds the new one. + +```bash +pip install 'qdrant-client>=1.9' # Qdrant only; pgvector needs no extra package +make scale-up SERVICE=qdrant # or pgvector (reuses POSTGRES_USER / POSTGRES_PASSWORD from .env, port 5433) +# .env: VECTOR_STORE=qdrant QDRANT_URL=http://localhost:6333 (pgvector: VECTOR_STORE=pgvector) +make sync ADMIN= FORCE=1 +make status +``` + +Settings: `QDRANT_URL`, `QDRANT_API_KEY` (empty means no authentication: fine on loopback, set one before the port is +reachable from anywhere else), `QDRANT_PREFIX` (prepended to collection names so installs can share a server), +`PGVECTOR_URI` (default: derived from `POSTGRES_*` for the compose profile), `PGVECTOR_SCHEMA` (default `public`). + +What every store guarantees, and the tests that pin it (`tests/test_vector_stores.py`, run against Chroma, embedded +Qdrant, a real Qdrant server and a real pgvector server): + +* **The clearance guard is in the base class.** Each store filters by clearance on the server, and then + `VectorStore.query()` drops any hit with a missing or invalid clearance, one the caller may not read, or another + workspace's. A broken filter in an adapter cannot leak a chunk. A test with a deliberately leaky adapter proves it. +* **Similarity is cosine on Chroma's scale.** The relevance gates are calibrated on it, so the values are compared, not + just the order (within 1e-4 of Chroma on the same corpus). `make eval` passes through each real adapter with the same + numbers as Chroma (answer rate and retrieval recall 0.9773, abstention accuracy 1.0). +* **A publish is atomic.** Qdrant builds a new collection and moves an alias in one request; pgvector replaces the + workspace's rows in one transaction. A publish that fails part-way leaves the old index serving. +* **No unfiltered fallback.** If the filter fails, dense retrieval becomes unavailable (labelled) instead of asking for + everything and trusting the guard. Chroma keeps its old fallback for old collections, behind the same guard. + +Measured on this machine, 384-dimensional random vectors, top 5 with a clearance filter, servers on loopback (one run; +random vectors are the worst case for an HNSW build): + +| Chunks | Store | Publish | Query median (p95) | +|---:|---|---:|---:| +| 4,000 | Chroma (embedded) | 1.64 s | 8.5 ms (11.5) | +| 4,000 | Qdrant | 0.76 s | 6.9 ms (11.9) | +| 4,000 | pgvector (exact) | 0.88 s | 3.5 ms (6.7) | +| 20,000 | Chroma (embedded) | 11.9 s | 31.9 ms (37.0) | +| 20,000 | Qdrant | 3.0 s | 9.6 ms (13.0) | +| 20,000 | pgvector (exact) | 4.3 s | 12.0 ms (15.4) | + +A whole query takes about 240 ms, and embedding the question is 4.5 ms of it: at the sizes this project is built for the +stores differ by single milliseconds. Choose one to **share an index**, not for speed. + +pgvector creates no ANN index: the column has no fixed dimension (the model can change) and HNSW needs one. The exact scan +grows linearly; beyond a few hundred thousand chunks per workspace, add an HNSW index on the cast column yourself. + +## Chat model: vLLM or any OpenAI-compatible server + +`LLM_PROVIDER=openai` talks to any server that implements `POST /v1/chat/completions` (vLLM, LM Studio, llama.cpp). + +| Setting | Meaning | +|---|---| +| `LOCAL_LLM_URL` | Base URL including `/v1`, for example `http://localhost:8000/v1` | +| `LOCAL_LLM_MODEL` | The model id the server serves (for vLLM, the `--model` value). Falls back to `OLLAMA_MODEL` | +| `LOCAL_LLM_API_KEY` | Sent as `Authorization: Bearer ...`. A secret: never printed, never in `repr()` | +| `LOCAL_LLM_TEMPLATE_KWARGS` | JSON merged in as `chat_template_kwargs`. The default `{"enable_thinking": false}` turns off Qwen3's thinking, which vLLM leaves **on**. Set it empty for a server that rejects unknown fields | + +Any `` block left in the text is removed (so is one that never closed), and a separate `reasoning` field is +ignored: the graders and the reader see only the answer. `make status` probes `GET /v1/models` and says whether the server +is unreachable, refused the key, or does not serve the model. If the server fails during a question, the answer falls +back to the labelled extractive one, as with Ollama. Tested against a stub that speaks the vLLM API over a Unix socket +(`tests/test_openai_server.py`). + +`make scale-up SERVICE=vllm` starts the official image with `VLLM_MODEL` (default `Qwen/Qwen3-8B`) and `VLLM_API_KEY` +(empty: no authentication, loopback only). It needs an NVIDIA GPU and the NVIDIA Container Toolkit, so it cannot run on a +Mac, and **that profile has only had its configuration validated, never run**. The first start downloads the model from +Hugging Face. Unlike the ONNX models it is not SHA-256 pinned, and the host is not air-gapped until the cache volume is +filled and `HF_HUB_OFFLINE=1` is set for the container. + +## Shared semantic cache + +`CACHE_BACKEND=postgres` keeps the cache table in PostgreSQL instead of a local SQLite file, so several UI processes read +and write one cache. Set `CACHE_DB_URI` to a database (it can be the pgvector database or any other). + +* **Sealing is unchanged.** Question, answer, citations and query vector are AES-GCM sealed with the row identity as + additional data, so the database operator learns none of them and a row moved to another clearance fails to open. The + columns that are not sealed are the same as in SQLite: workspace, clearance, index version, review flag, approver, time. +* **Every process needs the same data key and the same access key.** On one host, copy the key files or point + `DATA_KEY_PATH` / `ACCESS_KEY_PATH` at shared ones. On separate hosts `KEY_STORE=auto` would create a *different* key + per host (each OS keychain is local), so cache rows would silently never open and session tokens would not verify across + hosts: use `KEY_STORE=file` with one distributed key file. Envelope encryption through a KMS is an interface and a stub only. +* **L1 is off.** The in-process answer cache cannot be evicted by another process's prune or delete, so every lookup + reads the shared table. Semantic search still uses the in-memory vector index, re-synced with the table on every lookup. +* **The index version is part of the key.** It is a hash of the workspace's manifest in `bm25_indices/`, so processes + that answer the same workspace must hold the same manifest (next section). + +## What does not scale + +The options above remove the vector store, the chat model and the cache from the list of single-host components. These +remain **local files**, and a deployment over several hosts must put each on shared storage or keep one writer: + +* the keyword index (`bm25_indices/`, SQLite FTS5) and the manifests beside it, which the cache version depends on; +* `uploads/` and `knowledge_base/`, the source documents; +* the audit log (`data/audit.db`, a SQLite hash chain), its anchors, and the access policy (`access_policy.json`); +* the keys. + +Several UI **processes on one host** (behind a reverse proxy, for example) work with all of these as they are: SQLite is +safe across processes on one machine. SQLite over a network file system is not, which is the case the shared cache removes +for the cache alone. Checkpoints can already be PostgreSQL (`DB_URI`, `make up`). + +`make backup` covers the audit log, anchors, checkpoints and policy of this host. It does **not** back up Qdrant, pgvector +or a shared cache database; those are yours to back up (`docker volume` snapshots, `pg_dump`), and any of them can be rebuilt +from the documents with `make sync FORCE=1`. + +## Compose profiles + +| Profile | Image | Port (loopback) | Volume | Authentication | +|---|---|---|---|---| +| `qdrant` | `qdrant/qdrant:v1.20.0` | `QDRANT_PORT` (6333) | `qdrant_data` | `QDRANT_API_KEY`, empty = none | +| `pgvector` | `pgvector/pgvector:pg16` | `PGVECTOR_PORT` (5433) | `pgvector_data` | `POSTGRES_PASSWORD` from `.env` | +| `vllm` | `vllm/vllm-openai:latest` | `VLLM_PORT` (8000) | `vllm_cache` | `VLLM_API_KEY`, empty = none | +| `caddy` | `caddy:2-alpine` | `RAG_HTTPS_PORT` (8443) | `caddy_data` | see [lan-tls.md](lan-tls.md) | + +`pgvector/pgvector:pg16` and `vllm/vllm-openai:latest` are moving tags: pin a digest for a deployment you keep. Nothing +here is exposed beyond loopback by default; put a TLS proxy or a private network in front before reaching a service from +another host. diff --git a/tests/test_scaling_config.py b/tests/test_scaling_config.py new file mode 100644 index 0000000..a61ae87 --- /dev/null +++ b/tests/test_scaling_config.py @@ -0,0 +1,112 @@ +"""The scaling options are off by default, loopback-only, validated, documented, and cannot leak into tests.""" +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +from rag_lab.config import settings + +ROOT = Path(__file__).resolve().parent.parent +COMPOSE = yaml.safe_load((ROOT / "docker-compose.yml").read_text(encoding="utf-8")) +SERVICES = COMPOSE["services"] +SCALING = ("qdrant", "pgvector", "vllm") +MAKEFILE = (ROOT / "Makefile").read_text(encoding="utf-8") + + +def test_make_up_starts_only_the_checkpoint_database(): + assert {name for name, svc in SERVICES.items() if "profiles" not in svc} == {"postgres"} + up = MAKEFILE.split("\nup:")[1].split("\n\n")[0] + assert not any(name in up for name in (*SCALING, "caddy")) and "--profile" not in up + + +@pytest.mark.parametrize("name", [*SCALING, "caddy"]) +def test_each_option_has_exactly_one_profile_named_after_it(name): + assert SERVICES[name]["profiles"] == [name] + + +@pytest.mark.parametrize("name", SCALING) +def test_scaling_services_listen_on_loopback_only(name): + ports = SERVICES[name]["ports"] + assert ports and all(str(p).startswith("127.0.0.1:") for p in ports), ports + + +def test_images_are_named_and_the_moving_ones_are_documented(): + assert SERVICES["qdrant"]["image"] == "qdrant/qdrant:v1.20.0", "pinned: the client and the server move together" + docs = (ROOT / "docs" / "scaling.md").read_text(encoding="utf-8") + for floating in ("pgvector/pgvector:pg16", "vllm/vllm-openai:latest"): + assert SERVICES[next(n for n, s in SERVICES.items() if s["image"] == floating)]["image"] == floating + assert floating in docs and "pin a digest" in docs + + +def test_no_secret_is_written_into_compose(): + text = (ROOT / "docker-compose.yml").read_text(encoding="utf-8") + assert "rag_password" not in text + assert SERVICES["pgvector"]["environment"]["POSTGRES_PASSWORD"].startswith("${POSTGRES_PASSWORD:?"), "no default password" + assert SERVICES["qdrant"]["environment"]["QDRANT__SERVICE__API_KEY"] == "${QDRANT_API_KEY:-}" + assert SERVICES["vllm"]["environment"]["VLLM_API_KEY"] == "${VLLM_API_KEY:-}" + + +def test_vllm_asks_for_a_gpu_because_it_cannot_run_without_one(): + devices = SERVICES["vllm"]["deploy"]["resources"]["reservations"]["devices"] + assert devices[0]["driver"] == "nvidia" and "gpu" in devices[0]["capabilities"] + docs = (ROOT / "docs" / "scaling.md").read_text(encoding="utf-8") + assert "never run" in docs, "the docs must say this profile was validated, not run" + + +@pytest.mark.skipif(shutil.which("docker") is None, reason="docker is not installed") +@pytest.mark.parametrize("profile", [*SCALING, "caddy"]) +def test_every_profile_is_a_valid_compose_configuration(profile): + env = {**os.environ, "POSTGRES_PASSWORD": "not-a-real-password"} + done = subprocess.run(["docker", "compose", "--profile", profile, "config", "--quiet"], cwd=ROOT, env=env, + capture_output=True, text=True) + if done.returncode != 0 and "Cannot connect to the Docker daemon" in done.stderr: + pytest.skip("the docker daemon is not running") + assert done.returncode == 0, done.stderr + + +def test_the_make_targets_exist_are_listed_and_check_their_argument(): + assert re.search(r"^\.PHONY:.*\bscale-up\b.*\bscale-down\b", MAKEFILE, re.MULTILINE) + help_text = MAKEFILE.split("help:")[1].split("status:")[0] + assert "make scale-up SERVICE=" in help_text and "make scale-down SERVICE=" in help_text + for target in ("scale-up", "scale-down"): + recipe = MAKEFILE.split(f"\n{target}:")[1].split("\n\n")[0] + assert "qdrant|pgvector|vllm" in recipe and "--profile $(SERVICE)" in recipe + + +# --- defaults and isolation ---------------------------------------------------------------------------------------- + +def test_every_option_is_off_by_default(): + assert settings.vector_store == "chroma" and settings.cache_backend == "sqlite" + assert settings.llm_provider == "deterministic", "tests run on the in-process synthesizer" + assert not settings.qdrant_url and not settings.pgvector_uri and not settings.cache_db_uri + + +def test_the_eval_redirect_cannot_be_hijacked_by_an_env_file(): + source = (ROOT / "rag_lab" / "ops" / "eval.py").read_text(encoding="utf-8") + assert 'os.environ["CACHE_BACKEND"] = "sqlite"' in source and 'os.environ["CACHE_DB_URI"] = ""' in source + assert 'os.environ.setdefault("QDRANT_PREFIX", "eval_")' in source and 'os.environ.setdefault("PGVECTOR_SCHEMA", "eval")' in source + + +# --- documentation does not drift from the code ------------------------------------------------------------------------------ + +SETTING_VARS = ("VECTOR_STORE", "QDRANT_URL", "QDRANT_API_KEY", "QDRANT_PREFIX", "PGVECTOR_URI", "PGVECTOR_SCHEMA", + "CACHE_BACKEND", "CACHE_DB_URI", "LOCAL_LLM_URL", "LOCAL_LLM_MODEL", "LOCAL_LLM_API_KEY", + "LOCAL_LLM_TEMPLATE_KWARGS") + + +@pytest.mark.parametrize("variable", SETTING_VARS) +def test_every_scaling_variable_is_read_by_the_code_and_documented_twice(variable): + assert f'"{variable}"' in (ROOT / "rag_lab" / "config.py").read_text(encoding="utf-8") + assert f"{variable}" in (ROOT / ".env.example").read_text(encoding="utf-8") + assert variable in (ROOT / "docs" / "scaling.md").read_text(encoding="utf-8") + + +def test_the_docs_state_what_does_not_scale(): + docs = (ROOT / "docs" / "scaling.md").read_text(encoding="utf-8") + for local_file in ("bm25_indices/", "audit.db", "access_policy.json"): + assert local_file in docs + assert "same data key" in docs and "KEY_STORE=file" in docs diff --git a/tests/test_vector_stores.py b/tests/test_vector_stores.py index eab2052..a6fde66 100644 --- a/tests/test_vector_stores.py +++ b/tests/test_vector_stores.py @@ -80,6 +80,8 @@ def store(request, tmp_path, monkeypatch): if kind == "qdrant-local": adapter = QdrantStore(":memory:", "", prefix) else: + if os.getenv("RAG_OFFLINE") == "1": + pytest.skip("the offline run blocks TCP connections made through Python sockets, a Qdrant server included") url = os.getenv("QDRANT_TEST_URL") if not url: _unavailable("QDRANT_TEST_URL is not set (start one with `make scale-up SERVICE=qdrant`)", "REQUIRE_QDRANT") From dde1ef6b500120794bdf2a636c192a8b4a678102 Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 21:00:55 +0200 Subject: [PATCH 6/8] fix(security): connection strings and API keys never appear in a dict key The shared-cache pool registry and the vector-store instance cache were keyed by the connection string / API key. A dict key is printed by tracebacks and by pytest's assertion output, and an order-dependent test did exactly that. Both are keyed by a digest now, with tests that no key contains the secret, and the order-dependent test checks behaviour (the SQLite backend never opens a pool) instead of global state. --- rag_lab/retrieval/semantic_cache.py | 11 ++++++---- rag_lab/storage/vector_store.py | 10 +++++++-- tests/test_shared_cache.py | 34 +++++++++++++++++++++++++++-- tests/test_vector_stores.py | 11 ++++++++++ 4 files changed, 58 insertions(+), 8 deletions(-) diff --git a/rag_lab/retrieval/semantic_cache.py b/rag_lab/retrieval/semantic_cache.py index 16cefe2..9e3bace 100644 --- a/rag_lab/retrieval/semantic_cache.py +++ b/rag_lab/retrieval/semantic_cache.py @@ -269,19 +269,22 @@ def _pg_pool() -> Any: uri = settings.cache_db_uri if not uri: raise RuntimeError("CACHE_BACKEND=postgres needs CACHE_DB_URI (a PostgreSQL connection string).") + # Keyed by a digest, never by the URI: the URI carries the password, and a dict key is printed by tracebacks and by + # test assertion output. + key = hashlib.sha256(uri.encode("utf-8")).hexdigest()[:16] with _PG_LOCK: - pool = _PG_POOLS.get(uri) + pool = _PG_POOLS.get(key) if pool is None: pool = ConnectionPool(conninfo=uri, min_size=1, max_size=max(2, settings.checkpoint_pool_size), open=True, timeout=15, kwargs={"autocommit": False}) atexit.register(pool.close) # its worker threads must be stopped before the interpreter finalises - _PG_POOLS[uri] = pool - if uri not in _PG_READY: + _PG_POOLS[key] = pool + if key not in _PG_READY: with pool.connection() as raw: raw.execute("SELECT pg_advisory_xact_lock(727274)") # two processes starting together create it once for statement in _PG_DDL: raw.execute(statement) - _PG_READY.add(uri) + _PG_READY.add(key) return pool diff --git a/rag_lab/storage/vector_store.py b/rag_lab/storage/vector_store.py index e77a21a..5faab71 100644 --- a/rag_lab/storage/vector_store.py +++ b/rag_lab/storage/vector_store.py @@ -19,6 +19,7 @@ """ from __future__ import annotations +import hashlib import logging from abc import ABC, abstractmethod from collections.abc import Iterator, Sequence @@ -260,13 +261,18 @@ def _search(self, workspace, vector, n, allowed_clearances) -> list[Hit]: _INSTANCES: dict[tuple, VectorStore] = {} +def _digest(secret: str) -> str: + """Names a connection string or key in a dict key without containing it (tracebacks and test output print keys).""" + return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16] + + def get_vector_store() -> VectorStore: """The store selected by VECTOR_STORE, read when called so a settings override applies.""" kind = settings.vector_store if kind == "chroma": return ChromaStore() if kind == "qdrant": - key: tuple = ("qdrant", settings.qdrant_url, settings.qdrant_api_key, settings.qdrant_prefix) + key: tuple = ("qdrant", settings.qdrant_url, _digest(settings.qdrant_api_key), settings.qdrant_prefix) if key not in _INSTANCES: try: from rag_lab.storage.qdrant_store import QdrantStore @@ -278,7 +284,7 @@ def get_vector_store() -> VectorStore: _INSTANCES[key] = QdrantStore(settings.qdrant_url, settings.qdrant_api_key, settings.qdrant_prefix) return _INSTANCES[key] if kind == "pgvector": - key = ("pgvector", settings.pgvector_uri, settings.pgvector_schema) + key = ("pgvector", _digest(settings.pgvector_uri), settings.pgvector_schema) if key not in _INSTANCES: from rag_lab.storage.pgvector_store import PgvectorStore diff --git a/tests/test_shared_cache.py b/tests/test_shared_cache.py index e57d6de..812fabe 100644 --- a/tests/test_shared_cache.py +++ b/tests/test_shared_cache.py @@ -25,8 +25,38 @@ def test_sqlite_only_sql_stays_in_the_sqlite_constants(): assert statements and not [s for s in statements if "%" in s], "a literal % would be read as a placeholder" -def test_the_postgres_pool_is_not_created_unless_the_backend_is_selected(): - assert settings.cache_backend == "sqlite" and not sc._PG_POOLS +def test_the_postgres_pool_is_never_touched_while_the_backend_is_sqlite(tmp_path, monkeypatch): + def boom(): + raise AssertionError("the SQLite backend must not open a PostgreSQL pool") + + monkeypatch.setattr(sc, "_pg_pool", boom) + assert settings.cache_backend == "sqlite" + sc.store_cache("ops", "Where is the runbook?", "public", "in the wiki", db_path=str(tmp_path / "c.db"), compute_vector=False) + assert sc.lookup_cache("ops", "Where is the runbook?", "public", db_path=str(tmp_path / "c.db"))["hit"] + + +def test_the_pool_registry_never_holds_a_connection_string(monkeypatch): + """A dict key is printed by tracebacks and by pytest's assertion output: it must not carry a password.""" + uri = "postgresql://user:hunter2-not-a-real-password@localhost:5432/db" + monkeypatch.setattr(settings, "cache_backend", "postgres") + monkeypatch.setattr(settings, "cache_db_uri", uri) + + class FakePool: + def __init__(self, **kw): + self.closed = False + + def connection(self): + raise RuntimeError("never reached: the test only inspects the registry key") + + def close(self): + pass + + monkeypatch.setattr("psycopg_pool.ConnectionPool", FakePool) + monkeypatch.setattr(sc, "_PG_POOLS", {}) + monkeypatch.setattr(sc, "_PG_READY", set()) + with pytest.raises(RuntimeError, match="never reached"): + sc._pg_pool() + assert sc._PG_POOLS and not any("hunter2" in k or "postgresql" in k for k in sc._PG_POOLS) def test_a_missing_uri_is_a_labelled_error(monkeypatch): diff --git a/tests/test_vector_stores.py b/tests/test_vector_stores.py index a6fde66..c1f7bf2 100644 --- a/tests/test_vector_stores.py +++ b/tests/test_vector_stores.py @@ -375,3 +375,14 @@ def test_the_pgvector_pool_is_closed_at_exit_so_the_interpreter_finishes_cleanly with psycopg.connect(uri, autocommit=True) as conn: conn.execute(sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(sql.Identifier(schema))) assert done.returncode == 0 and "PythonFinalizationError" not in done.stderr, done.stderr[-400:] + + +def test_the_instance_cache_never_holds_a_key_or_a_connection_string(monkeypatch): + """Its keys are printed by tracebacks and by pytest's assertion output.""" + pytest.importorskip("qdrant_client") + monkeypatch.setattr(settings, "vector_store", "qdrant") + monkeypatch.setattr(settings, "qdrant_url", ":memory:") + monkeypatch.setattr(settings, "qdrant_api_key", "hunter2-not-a-real-key") + monkeypatch.setattr(vs, "_INSTANCES", {}) + assert get_vector_store().name == "qdrant" + assert not any("hunter2" in str(part) for key in vs._INSTANCES for part in key) From ba7ec17d7345fc893c2495ef0e7e1f0593749e15 Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 21:04:54 +0200 Subject: [PATCH 7/8] test: the 'every module imports' check knows the optional qdrant module A fresh clone without qdrant-client failed it. Optional modules may fail to import only because exactly that package is missing; anything else still fails the test. --- tests/test_packaging.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_packaging.py b/tests/test_packaging.py index 6779bfd..bf41213 100644 --- a/tests/test_packaging.py +++ b/tests/test_packaging.py @@ -77,12 +77,21 @@ def test_package_modules_import_and_ship_their_data(): assert {"common_passwords.txt", "injection_patterns.tsv"} <= set(data["tool"]["setuptools"]["package-data"]["rag_lab.security"]) assert "data/*.json" in data["tool"]["setuptools"]["package-data"]["rag_lab.lang"] + # Modules that need an optional package (an extra in pyproject.toml): they must import where it is installed and may + # fail to import, with exactly that package missing, where it is not. + optional = {"rag_lab.storage.qdrant_store": "qdrant_client"} + assert set(optional.values()) <= {"qdrant_client"} and "qdrant" in data["project"]["optional-dependencies"] + names = [m.name for m in pkgutil.walk_packages(rag_lab.__path__, "rag_lab.")] assert len(names) >= 30 for name in names: # The Streamlit page runs only inside a Streamlit session (see the root ui.py launcher). - if name != "rag_lab.ui.streamlit_app": + if name == "rag_lab.ui.streamlit_app": + continue + try: assert importlib.import_module(name) is not None + except ModuleNotFoundError as exc: + assert optional.get(name) == exc.name, f"{name} failed to import: {exc}" def test_no_flat_modules_are_left_at_the_repository_root(): From 59ad049134e91de9d1bb2de3777a043adff93802 Mon Sep 17 00:00:00 2001 From: IndieGoWeb Ltd Date: Wed, 4 Nov 2026 21:12:42 +0200 Subject: [PATCH 8/8] fix(ops): status asks the vector store and the shared cache; factory-reset purges them; eval namespaces are forced make status no longer reports healthy for a dead backend: a Vector store line (name and location, never a credential) and a Semantic cache line that actually asks the shared PostgreSQL, both [!!] and exit 1 when they do not answer. make factory-reset purges the workspaces of this install's Qdrant prefix / pgvector schema and the shared cache table before deleting the local folders, and prints NOT PURGED (exit 1) for anything it could not reach or could not empty, so its 'zero-trace' label is true. The eval redirect forces QDRANT_PREFIX=eval_ and PGVECTOR_SCHEMA=eval, so an exported real prefix or schema cannot receive the eval corpus. A test no longer passes the pgvector URI on argv. docs/scaling.md now states the trust boundary (Qdrant and pgvector hold every chunk in the clear; RBAC is this application's, not the store's) and what was checked live: Ollama's /v1 accepts chat_template_kwargs and ignores it (Qwen3 still thinks and returns it in a separate reasoning field, which the provider ignores), so the answer stays clean. --- Makefile | 3 +- docs/scaling.md | 25 +++- rag_lab/ops/eval.py | 7 +- rag_lab/ops/purge_external.py | 65 ++++++++++ rag_lab/retrieval/semantic_cache.py | 13 ++ rag_lab/storage/pgvector_store.py | 6 + rag_lab/storage/qdrant_store.py | 7 ++ rag_lab/storage/vector_store.py | 27 +++++ rag_lab/storage/workspace_manager.py | 21 +++- tests/test_external_stores_ops.py | 173 +++++++++++++++++++++++++++ tests/test_scaling_config.py | 3 +- tests/test_vector_stores.py | 22 +++- 12 files changed, 356 insertions(+), 16 deletions(-) create mode 100644 rag_lab/ops/purge_external.py create mode 100644 tests/test_external_stores_ops.py diff --git a/Makefile b/Makefile index 63f9428..2216661 100644 --- a/Makefile +++ b/Makefile @@ -400,13 +400,14 @@ reset-users: # 20. Complete zero-trace factory reset (wipes policy, data, audit, checkpoints, keys, vectors, uploads, and caches) factory-reset: @if [ "$(CONFIRM)" != "yes" ] && [ "$(YES)" != "1" ] && [ "$(YES)" != "yes" ]; then \ - printf "WARNING: This permanently deletes all user accounts, audit logs, checkpoints, encryption keys, and indexed vectors.\nType 'yes' to proceed: "; \ + printf "WARNING: This permanently deletes all user accounts, audit logs, checkpoints, encryption keys, and indexed vectors, including whatever the configured Qdrant, pgvector or shared-cache server holds for this install.\nType 'yes' to proceed: "; \ read ans; \ if [ "$$ans" != "yes" ]; then \ echo "Factory reset cancelled."; \ exit 1; \ fi; \ fi + ./venv/bin/python -m rag_lab.ops.purge_external || echo "Warning: a configured external store was NOT purged (see above); the documents or cached answers are still on that server." ./venv/bin/python -m rag_lab.security.keystore forget || echo "Warning: keychain entries could not be removed; delete the local-rag-lab entries in your keychain app by hand." rm -f access_policy.json access_policy.json.*.bak access_policy.json.lock rm -rf data diff --git a/docs/scaling.md b/docs/scaling.md index 04be85a..6d9108e 100644 --- a/docs/scaling.md +++ b/docs/scaling.md @@ -65,6 +65,23 @@ stores differ by single milliseconds. Choose one to **share an index**, not for pgvector creates no ANN index: the column has no fixed dimension (the model can change) and HNSW needs one. The exact scan grows linearly; beyond a few hundred thousand chunks per workspace, add an HNSW index on the cast column yourself. +### Trust boundary: the stores hold your documents in the clear + +**The chunks are not encrypted in Qdrant or pgvector.** The text of every chunk, at every clearance level, and its +clearance tag sit in Qdrant's payloads or in pgvector's `document` column, readable by anyone who can read that server +or its disk: a database administrator, a backup, another application with access. Role-based access control is enforced +**by this application** (the clearance filter and the guard), not by the store. Embedded Chroma on the same host has the +same property; what changes is that its files are local. The shared semantic cache is different: it is sealed, as above. + +So treat a Qdrant or pgvector server as holding your most sensitive documents: keep it on loopback or a private network, +set `QDRANT_API_KEY` or a pgvector password, put TLS in front of anything that crosses a network, encrypt the disk and its +backups, and give nobody who must not read `secret` documents access to it. + +`make factory-reset` empties the stores this install is configured for (every workspace of the Qdrant prefix or pgvector +schema, and the shared cache table) before it deletes the local folders. If a server cannot be reached it prints +`NOT PURGED` for it and says the documents are still there; it never reports a wipe that did not happen. Anything else on a +shared server (other installs, other prefixes) is untouched. + ## Chat model: vLLM or any OpenAI-compatible server `LLM_PROVIDER=openai` talks to any server that implements `POST /v1/chat/completions` (vLLM, LM Studio, llama.cpp). @@ -76,9 +93,13 @@ grows linearly; beyond a few hundred thousand chunks per workspace, add an HNSW | `LOCAL_LLM_API_KEY` | Sent as `Authorization: Bearer ...`. A secret: never printed, never in `repr()` | | `LOCAL_LLM_TEMPLATE_KWARGS` | JSON merged in as `chat_template_kwargs`. The default `{"enable_thinking": false}` turns off Qwen3's thinking, which vLLM leaves **on**. Set it empty for a server that rejects unknown fields | -Any `` block left in the text is removed (so is one that never closed), and a separate `reasoning` field is +What a server does with that field varies, and only some of it has been checked: vLLM honours it according to its +documentation, which I have not run live; **Ollama's `/v1` endpoint accepts it and ignores it** (checked live with +`qwen3:8b`: it still thinks, about 650 characters for a one-line question, and returns the thinking in a separate +`reasoning` field). That is why the provider ignores `reasoning` fields rather than trusting the setting: the answer +stays clean either way. For Ollama, use `LLM_PROVIDER=ollama`, which turns thinking off natively. Any `` block left in the text is removed (so is one that never closed), and a separate `reasoning` field is ignored: the graders and the reader see only the answer. `make status` probes `GET /v1/models` and says whether the server -is unreachable, refused the key, or does not serve the model. If the server fails during a question, the answer falls +is unreachable, refused the key, or does not serve the model; it also reports the vector store and the shared cache (name and location, never a credential) and exits 1 if either does not answer. If the server fails during a question, the answer falls back to the labelled extractive one, as with Ollama. Tested against a stub that speaks the vLLM API over a Unix socket (`tests/test_openai_server.py`). diff --git a/rag_lab/ops/eval.py b/rag_lab/ops/eval.py index 0572b4a..a25bb5f 100644 --- a/rag_lab/ops/eval.py +++ b/rag_lab/ops/eval.py @@ -64,10 +64,11 @@ def redirect_stores(out_dir: Path, llm_mode: bool) -> None: os.environ["AUDIT_ANCHOR_COPY_DIR"] = "" # The evaluation measures the embedded defaults; a .env that enables a scaling option must not redirect it. # To evaluate an adapter on purpose, say so in the shell: `VECTOR_STORE=qdrant make eval` (docs/scaling.md). Even - # then it gets its own namespace, never the real collections or schema. + # then it gets its own namespace, forced here: an exported real prefix or schema must never receive the eval corpus, + # whose workspaces (engineering, finance, hr...) would overwrite real ones. os.environ.setdefault("VECTOR_STORE", "chroma") - os.environ.setdefault("QDRANT_PREFIX", "eval_") - os.environ.setdefault("PGVECTOR_SCHEMA", "eval") + os.environ["QDRANT_PREFIX"] = "eval_" + os.environ["PGVECTOR_SCHEMA"] = "eval" os.environ["CACHE_BACKEND"] = "sqlite" os.environ["CACHE_DB_URI"] = "" os.environ["DB_URI"] = f"sqlite:///{out_dir / 'unused_checkpoints.db'}" diff --git a/rag_lab/ops/purge_external.py b/rag_lab/ops/purge_external.py new file mode 100644 index 0000000..be6774d --- /dev/null +++ b/rag_lab/ops/purge_external.py @@ -0,0 +1,65 @@ +"""Empty the external stores this install is configured to use. Part of `make factory-reset`. + +Under VECTOR_STORE=qdrant|pgvector the indexed documents live on a server, and under CACHE_BACKEND=postgres the cached +answers do; deleting the local folders does not touch them, so a "zero-trace" reset must. Scope: every workspace in this +install's namespace (the Qdrant prefix, the pgvector schema) and the whole shared cache table. If a store cannot be +reached the line says NOT PURGED, loudly, and the command exits non-zero: it never claims a wipe that did not happen. +""" +from __future__ import annotations + +import sys + +from rag_lab.config import settings + + +def purge() -> tuple[list[str], bool]: + """(report lines, everything configured was purged).""" + lines: list[str] = [] + complete = True + if settings.vector_store != "chroma": + from rag_lab.storage.vector_store import get_vector_store, vector_store_status + + healthy, status = vector_store_status() + if not healthy: + lines.append(f"NOT PURGED: {status}. The indexed documents are still on that server.") + complete = False + else: + store = get_vector_store() + names = sorted(store.workspaces()) + for name in names: + store.delete(name) + left = store.workspaces() + if left: + lines.append(f"NOT PURGED: {len(left)} workspace(s) remain in {store.describe()}.") + complete = False + else: + lines.append(f"Purged {len(names)} workspace(s) from {store.describe()}.") + if settings.cache_backend == "postgres": + from rag_lab.retrieval.semantic_cache import cache_backend_status, cache_stats, clear_cache + + healthy, where = cache_backend_status() + if not healthy: + lines.append(f"NOT PURGED: {where}. The cached answers are still there.") + complete = False + else: + clear_cache() + remaining = cache_stats()["total_entries"] + if remaining: + lines.append(f"NOT PURGED: {remaining} cached answer(s) remain in the {where}.") + complete = False + else: + lines.append(f"Emptied the {where}.") + if not lines: + lines.append("No external store is configured: nothing to purge beyond the local folders.") + return lines, complete + + +def main() -> int: + lines, complete = purge() + for line in lines: + print(line) + return 0 if complete else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/rag_lab/retrieval/semantic_cache.py b/rag_lab/retrieval/semantic_cache.py index 9e3bace..fd12401 100644 --- a/rag_lab/retrieval/semantic_cache.py +++ b/rag_lab/retrieval/semantic_cache.py @@ -674,6 +674,19 @@ def clear_cache(db_path: str | None = None) -> None: logger.debug("Error clearing cache: %s", exc) +def cache_backend_status() -> tuple[bool, str]: + """(healthy, where) for `make status`. The shared backend is actually asked; cache_stats() would answer 0 entries.""" + if not _shared(): + return True, f"SQLite {settings.cache_db_path}" + where = f"shared PostgreSQL {_db_identity(None).removeprefix('postgres://')}" + try: + with _open_db() as conn: + conn.execute("SELECT 1") + except Exception as exc: + return False, f"{where} is unreachable ({type(exc).__name__}): nothing is cached or served from it" + return True, where + + def cache_stats(db_path: str | None = None) -> dict[str, Any]: """Retrieve cache entry counts, workspace distribution, and on-disk size.""" path = Path(db_path or settings.cache_db_path) diff --git a/rag_lab/storage/pgvector_store.py b/rag_lab/storage/pgvector_store.py index 68b3eaa..0ecf189 100644 --- a/rag_lab/storage/pgvector_store.py +++ b/rag_lab/storage/pgvector_store.py @@ -54,6 +54,7 @@ def __init__(self, uri: str, schema: str = "public") -> None: if not _SCHEMA.fullmatch(schema): raise VectorStoreError(f"PGVECTOR_SCHEMA must be a plain lowercase identifier; got '{schema}'.") self._schema = schema + self._where = uri self._pool = ConnectionPool(conninfo=uri, min_size=1, max_size=max(2, settings.checkpoint_pool_size), open=True, timeout=15, kwargs={"autocommit": False}) # The pool's worker threads must be stopped before the interpreter finalises (Python 3.14 raises otherwise). @@ -68,6 +69,11 @@ def __init__(self, uri: str, schema: str = "public") -> None: "pgvector extension: use the `pgvector` compose profile (image pgvector/pgvector), not plain postgres." ) from exc + def describe(self) -> str: + from rag_lab.storage.vector_store import safe_location + + return f"pgvector at {safe_location(self._where)} (schema {self._schema})" + def close(self) -> None: self._pool.close() diff --git a/rag_lab/storage/qdrant_store.py b/rag_lab/storage/qdrant_store.py index b214f8f..6aa86a7 100644 --- a/rag_lab/storage/qdrant_store.py +++ b/rag_lab/storage/qdrant_store.py @@ -39,10 +39,17 @@ def __init__(self, url: str, api_key: str = "", prefix: str = "") -> None: if not url: raise VectorStoreError("VECTOR_STORE=qdrant needs QDRANT_URL (the compose profile listens on http://localhost:6333).") self._prefix = prefix + self._url = url self._client = QdrantClient(location=":memory:") if url == ":memory:" else QdrantClient( url=url, api_key=api_key or None, timeout=30 ) + def describe(self) -> str: + from rag_lab.storage.vector_store import safe_location + + where = "in-process memory" if self._url == ":memory:" else safe_location(self._url) + return f"Qdrant at {where}" + (f" (prefix '{self._prefix}')" if self._prefix else "") + # --- naming ------------------------------------------------------------------------------------------------ def _alias(self, workspace: str) -> str: return f"{self._prefix}ws_{workspace}" diff --git a/rag_lab/storage/vector_store.py b/rag_lab/storage/vector_store.py index 5faab71..c044eac 100644 --- a/rag_lab/storage/vector_store.py +++ b/rag_lab/storage/vector_store.py @@ -22,6 +22,7 @@ import hashlib import logging from abc import ABC, abstractmethod +from urllib.parse import urlsplit from collections.abc import Iterator, Sequence from dataclasses import dataclass from typing import Any @@ -69,9 +70,19 @@ def guard_hits(hits: Sequence[Hit], workspace: str, allowed_clearances: Sequence return kept +def safe_location(uri: str) -> str: + """host:port/path of a URL, without any user or password.""" + parts = urlsplit(uri) + return f"{parts.hostname or uri}{':' + str(parts.port) if parts.port else ''}{parts.path}" + + class VectorStore(ABC): name = "" + def describe(self) -> str: + """Where this store is, for status output. Never contains a credential.""" + return self.name + # --- what an adapter implements ----------------------------------------------------------------- @abstractmethod def workspaces(self) -> set[str]: @@ -138,6 +149,10 @@ def _similarity(distance: float, space: str) -> float: class ChromaStore(VectorStore): name = "chroma" BATCH = 256 + + def describe(self) -> str: + return f"embedded Chroma at {settings.chroma_path}" + # Cosine distance keeps similarity meaningful whichever embedding model produced the vectors. METADATA = {"hnsw:space": "cosine"} @@ -256,6 +271,18 @@ def _search(self, workspace, vector, n, allowed_clearances) -> list[Hit]: ] +def vector_store_status() -> tuple[bool, str]: + """(healthy, one line) for `make status`: which store, where, and whether it answers.""" + try: + store = get_vector_store() + except VectorStoreError as exc: + return False, f"{settings.vector_store}: {exc}" + try: + return True, f"{store.describe()}: {len(store.workspaces())} workspace(s)" + except Exception as exc: + return False, f"{store.describe()} is unreachable ({type(exc).__name__}): dense retrieval is unavailable" + + # --- The factory --------------------------------------------------------------------------------------------------- _INSTANCES: dict[tuple, VectorStore] = {} diff --git a/rag_lab/storage/workspace_manager.py b/rag_lab/storage/workspace_manager.py index 2bdfb30..63ad702 100644 --- a/rag_lab/storage/workspace_manager.py +++ b/rag_lab/storage/workspace_manager.py @@ -13,7 +13,7 @@ from rag_lab.storage.checkpointer import is_postgres_uri, get_sqlite_path from rag_lab.ingestion.parsers import SUPPORTED_EXTENSIONS from rag_lab.retrieval.sparse_search import delete_sparse_index, has_sparse_index -from rag_lab.storage.vector_store import WorkspaceNotIndexed, get_vector_store +from rag_lab.storage.vector_store import WorkspaceNotIndexed, get_vector_store, vector_store_status logger = logging.getLogger("workspace_manager") @@ -440,6 +440,11 @@ def print_status() -> int: print(f"[{marker}] Key storage : {text}") ok = ok and marker != "!!" + # 3c. Where the vectors live, and whether that place answers + store_ok, store_line = vector_store_status() + print(f"[{'OK' if store_ok else '!!'}] Vector store : {store_line}") + ok = ok and store_ok + # 4. Workspaces workspaces = list_workspaces() if not workspaces: @@ -489,11 +494,17 @@ def print_status() -> int: print("[--] Anchor copy : none (set AUDIT_ANCHOR_COPY_DIR to a second location so a rewound " "anchor file is detectable)") - # 7. Semantic cache + # 7. Semantic cache (SQLite file, or the shared PostgreSQL, which is asked rather than assumed) try: - from rag_lab.retrieval.semantic_cache import cache_stats - cs = cache_stats() - print(f"[OK] Semantic cache : {cs['total_entries']} entry(s) across {len(cs['workspaces'])} workspace(s)") + from rag_lab.retrieval.semantic_cache import cache_backend_status, cache_stats + healthy, where = cache_backend_status() + if not healthy: + print(f"[!!] Semantic cache : {where}") + ok = False + else: + cs = cache_stats() + note = f" [{where}]" if settings.cache_backend != "sqlite" else "" + print(f"[OK] Semantic cache : {cs['total_entries']} entry(s) across {len(cs['workspaces'])} workspace(s){note}") except Exception as exc: print(f"[!!] Semantic cache error: {exc}") ok = False diff --git a/tests/test_external_stores_ops.py b/tests/test_external_stores_ops.py new file mode 100644 index 0000000..1e6996e --- /dev/null +++ b/tests/test_external_stores_ops.py @@ -0,0 +1,173 @@ +"""Operating the optional stores: `make factory-reset` purges them, `make status` asks them, neither shows a credential.""" +from pathlib import Path + +import pytest + +from rag_lab.config import settings +from rag_lab.ops import purge_external +from rag_lab.retrieval import semantic_cache as sc +from rag_lab.storage import vector_store as vs +from rag_lab.storage.vector_store import ChromaStore, safe_location, vector_store_status + +ROOT = Path(__file__).resolve().parent.parent + + +class FakeStore(vs.VectorStore): + """A store in a dict, to see exactly what purge asks of it.""" + + name = "fake" + + def __init__(self, workspaces=("alpha", "beta"), broken=False, stubborn=False): + self.data = {w: [f"{w} chunk"] for w in workspaces} + self.broken, self.stubborn = broken, stubborn + + def describe(self): + return "Fake store at fake.example:6333" + + def workspaces(self): + if self.broken: + raise ConnectionError("down") + return set(self.data) + + def delete(self, workspace): + if not self.stubborn: + self.data.pop(workspace, None) + + def ensure(self, workspace): pass + def exists(self, workspace): return workspace in self.data + def count(self, workspace): return len(self.data[workspace]) + def collection_metadata(self, workspace): return {} + def read_vectors(self, workspace): return iter(()) + def publish(self, *a, **k): pass + def _search(self, *a, **k): return [] + + +@pytest.fixture +def external(monkeypatch): + """VECTOR_STORE=qdrant (a fake one) so that purge and status have something external to deal with.""" + store = FakeStore() + monkeypatch.setattr(settings, "vector_store", "qdrant") + monkeypatch.setattr(vs, "get_vector_store", lambda: store) + return store + + +# --- factory-reset ----------------------------------------------------------------------------------------------------- + +def test_with_only_local_stores_there_is_nothing_external_to_purge(): + lines, complete = purge_external.purge() + assert complete and lines == ["No external store is configured: nothing to purge beyond the local folders."] + assert purge_external.main() == 0 + + +def test_the_configured_vector_store_is_emptied_for_this_install(external, capsys): + lines, complete = purge_external.purge() + assert complete and external.data == {} + assert lines == ["Purged 2 workspace(s) from Fake store at fake.example:6333."] + assert purge_external.main() == 0 and "Purged 0 workspace(s)" in capsys.readouterr().out + + +def test_an_unreachable_store_is_reported_as_not_purged_and_fails(external, capsys): + external.broken = True + lines, complete = purge_external.purge() + assert not complete and lines[0].startswith("NOT PURGED") and "still on that server" in lines[0] + assert purge_external.main() == 1 and "NOT PURGED" in capsys.readouterr().out + + +def test_a_store_that_does_not_actually_delete_is_not_reported_as_purged(external): + external.stubborn = True + lines, complete = purge_external.purge() + assert not complete and "NOT PURGED: 2 workspace(s) remain" in lines[0] and "Purged" not in lines[0] + + +def test_the_shared_cache_is_emptied_and_an_unreachable_one_is_not_claimed(cache_backend): + if cache_backend != "postgres": + pytest.skip("this is about the shared cache") + sc.store_cache("ops", "Where is the runbook?", "public", "in the wiki", compute_vector=False) + sc.store_cache("hr", "How many leave days?", "public", "twenty-five", compute_vector=False) + assert sc.cache_stats()["total_entries"] == 2 + lines, complete = purge_external.purge() + assert complete and lines[0].startswith("Emptied the shared PostgreSQL") and sc.cache_stats()["total_entries"] == 0 + + +def test_an_unreachable_shared_cache_is_not_claimed_as_purged(monkeypatch): + monkeypatch.setattr(settings, "cache_backend", "postgres") + monkeypatch.setattr(sc, "cache_backend_status", lambda: (False, "shared PostgreSQL db:5432/rag is unreachable (OperationalError)")) + lines, complete = purge_external.purge() + assert not complete and lines[0].startswith("NOT PURGED") and "cached answers are still there" in lines[0] + + +def test_factory_reset_purges_the_external_stores_before_it_deletes_the_local_folders(): + recipe = (ROOT / "Makefile").read_text(encoding="utf-8").split("\nfactory-reset:")[1].split("\n\n")[0] + assert recipe.index("rag_lab.ops.purge_external") < recipe.index("rm -rf chroma_data") + assert "NOT purged" in recipe and "Qdrant, pgvector or shared-cache" in recipe + + +# --- make status ------------------------------------------------------------------------------------------------------------- + +def test_the_vector_store_line_names_the_store_and_the_count_without_a_credential(): + healthy, line = vector_store_status() + assert healthy and line.startswith("embedded Chroma at ") and "workspace(s)" in line + + +def test_a_dead_vector_store_is_reported_unhealthy(external): + external.broken = True + healthy, line = vector_store_status() + assert not healthy and "Fake store at fake.example:6333 is unreachable (ConnectionError)" in line + + +def test_a_misconfigured_store_is_reported_unhealthy(monkeypatch): + monkeypatch.setattr(settings, "vector_store", "milvus") + healthy, line = vector_store_status() + assert not healthy and "VECTOR_STORE must be one of" in line + + +@pytest.mark.filterwarnings("ignore:Failed to obtain server version") # the client probes a host that does not exist +def test_locations_never_carry_a_user_or_a_password(): + assert safe_location("postgresql://user:hunter2@db.internal:5433/rag_vectors") == "db.internal:5433/rag_vectors" + assert safe_location("http://localhost:6333") == "localhost:6333" + pytest.importorskip("qdrant_client") + from rag_lab.storage.qdrant_store import QdrantStore + + described = QdrantStore("https://qdrant.internal:6333", "hunter2-key", "t_").describe() + assert described == "Qdrant at qdrant.internal:6333 (prefix 't_')" and "hunter2" not in described + + +def test_the_sqlite_cache_reports_itself_healthy(): + healthy, where = sc.cache_backend_status() + assert healthy and where.startswith("SQLite ") + + +def test_status_exits_1_when_the_vector_store_is_down(rollback_index, capsys, monkeypatch): + from rag_lab import accel + from rag_lab.storage.workspace_manager import print_status + + monkeypatch.setattr(accel, "status", lambda: accel.AccelStatus("preferred", True, version="test")) + monkeypatch.setattr(vs, "vector_store_status", lambda: (False, "Qdrant at db:6333 is unreachable (ConnectError): dense retrieval is unavailable")) + monkeypatch.setattr("rag_lab.storage.workspace_manager.vector_store_status", vs.vector_store_status) + code = print_status() + out = capsys.readouterr().out + assert code == 1 and "[!!] Vector store : Qdrant at db:6333 is unreachable" in out + + +def test_status_exits_1_when_the_shared_cache_is_down(rollback_index, capsys, monkeypatch): + from rag_lab import accel + from rag_lab.storage.workspace_manager import print_status + + monkeypatch.setattr(accel, "status", lambda: accel.AccelStatus("preferred", True, version="test")) + monkeypatch.setattr(sc, "cache_backend_status", lambda: (False, "shared PostgreSQL db:5432/rag is unreachable (OperationalError): nothing is cached or served from it")) + code = print_status() + out = capsys.readouterr().out + assert code == 1 and "[!!] Semantic cache : shared PostgreSQL db:5432/rag is unreachable" in out + assert "[OK] Semantic cache" not in out + + +def test_status_is_healthy_by_default(rollback_index, capsys, monkeypatch): + from rag_lab import accel + from rag_lab.storage.workspace_manager import print_status + + monkeypatch.setattr(accel, "status", lambda: accel.AccelStatus("preferred", True, version="test")) + code = print_status() + out = capsys.readouterr().out + assert code == 0, out + assert "[OK] Vector store : embedded Chroma at" in out and "[OK] Semantic cache" in out + assert isinstance(vs.get_vector_store(), ChromaStore) diff --git a/tests/test_scaling_config.py b/tests/test_scaling_config.py index a61ae87..9a6c24e 100644 --- a/tests/test_scaling_config.py +++ b/tests/test_scaling_config.py @@ -88,7 +88,8 @@ def test_every_option_is_off_by_default(): def test_the_eval_redirect_cannot_be_hijacked_by_an_env_file(): source = (ROOT / "rag_lab" / "ops" / "eval.py").read_text(encoding="utf-8") assert 'os.environ["CACHE_BACKEND"] = "sqlite"' in source and 'os.environ["CACHE_DB_URI"] = ""' in source - assert 'os.environ.setdefault("QDRANT_PREFIX", "eval_")' in source and 'os.environ.setdefault("PGVECTOR_SCHEMA", "eval")' in source + assert 'os.environ["QDRANT_PREFIX"] = "eval_"' in source and 'os.environ["PGVECTOR_SCHEMA"] = "eval"' in source, \ + "forced, not defaulted: an exported real prefix or schema must not receive the eval corpus" # --- documentation does not drift from the code ------------------------------------------------------------------------------ diff --git a/tests/test_vector_stores.py b/tests/test_vector_stores.py index c1f7bf2..b2f2230 100644 --- a/tests/test_vector_stores.py +++ b/tests/test_vector_stores.py @@ -364,13 +364,14 @@ def test_the_pgvector_pool_is_closed_at_exit_so_the_interpreter_finishes_cleanly _unavailable("no pgvector server", "REQUIRE_PGVECTOR") schema = f"pytest_exit_{uuid.uuid4().hex[:6]}" code = ( - "import sys\n" + "import os, sys\n" "from rag_lab.storage.pgvector_store import PgvectorStore\n" - "PgvectorStore(sys.argv[1], sys.argv[2]).ensure('ws')\n" # then the interpreter exits: atexit must close the pool + # The URI travels in the environment, never on argv: a failing assertion prints argv, and the URI has a password. + "PgvectorStore(os.environ['RAG_EXIT_TEST_URI'], sys.argv[1]).ensure('ws')\n" # then the interpreter exits ) try: - done = subprocess.run([sys.executable, "-c", code, uri, schema], capture_output=True, text=True, - env=os.environ | {"PGVECTOR_URI": ""}) + done = subprocess.run([sys.executable, "-c", code, schema], capture_output=True, text=True, + env=os.environ | {"PGVECTOR_URI": "", "RAG_EXIT_TEST_URI": uri}) finally: with psycopg.connect(uri, autocommit=True) as conn: conn.execute(sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(sql.Identifier(schema))) @@ -386,3 +387,16 @@ def test_the_instance_cache_never_holds_a_key_or_a_connection_string(monkeypatch monkeypatch.setattr(vs, "_INSTANCES", {}) assert get_vector_store().name == "qdrant" assert not any("hunter2" in str(part) for key in vs._INSTANCES for part in key) + + +def test_factory_reset_purge_empties_a_real_store_of_this_installs_workspaces(store, monkeypatch): + """`make factory-reset` runs rag_lab.ops.purge_external; here against every adapter, real servers included.""" + from rag_lab.ops import purge_external + + publish(store, "alpha", seed=1) + publish(store, "beta", seed=2) + monkeypatch.setattr(settings, "vector_store", "qdrant") # anything but chroma: the external branch + monkeypatch.setattr(vs, "get_vector_store", lambda: store) + lines, complete = purge_external.purge() + assert complete and lines[0].startswith("Purged 2 workspace(s) from ") and "hunter2" not in lines[0] + assert store.workspaces() == set() and not store.exists("alpha")