From 0ff9d8dd21c38d4a40aaeaa6bc03ce61a256a65a Mon Sep 17 00:00:00 2001 From: Mahdi Bazrafshan Date: Mon, 10 Aug 2026 14:12:45 +0330 Subject: [PATCH] feat(storage): add app settings and model-scoped collections Why: - Different Embedding Models must not share Qdrant collections; Admin needs persisted role defaults and Experiment provenance. Changes: - Scope collection names by model id; add app_settings, embedding provenance columns, and queries.expansion_tree. Impact: - Process/query write only Model Corpus collections; legacy unscoped names are read-only for Admin labeling. Co-authored-by: Cursor --- src/storage/qdrant.py | 198 +++++++++++++++++++++++++++++++++--------- src/storage/sqlite.py | 135 ++++++++++++++++++++++++++-- 2 files changed, 284 insertions(+), 49 deletions(-) diff --git a/src/storage/qdrant.py b/src/storage/qdrant.py index 1c066f4..f291a03 100644 --- a/src/storage/qdrant.py +++ b/src/storage/qdrant.py @@ -1,75 +1,123 @@ """Qdrant vector storage layer. -One collection per chunking strategy. Handles collection creation, -vector upsert, and similarity search. +One collection per (Strategy, Embedding Model) — the Model Corpus. +Names are always `{strategy}__{model_id}_collection` (ADR-0021). +Legacy unscoped `{strategy}_collection` names are recognized for Admin +labeling only; process/query no longer write them. """ +from __future__ import annotations + import logging import uuid from typing import Any -from qdrant_client import QdrantClient from qdrant_client.models import ( Distance, FieldCondition, Filter, + MatchAny, MatchValue, PointIdsList, PointStruct, VectorParams, ) -from src.core.config import settings +from src.chunking.embedding_models import LEGACY_CLOUD_MODEL_ID from src.core.dependencies import get_qdrant_client from src.core.exceptions import QdrantError -from src.core.models import Chunk, ChunkMetadata, StrategyName, chunk_to_metadata +from src.core.models import Chunk, StrategyName, chunk_to_metadata logger = logging.getLogger(__name__) -# Embedding dimension for text-embedding-3-small +# Backward-compat alias — prefer EmbeddingModelSpec.dimension VECTOR_DIMENSION = 1536 -def collection_name(strategy: StrategyName | str) -> str: - """Convention: {strategy_name}_collection.""" +def _sanitize_model_id(model_id: str) -> str: + return model_id.replace(":", "-").replace("/", "-") + + +def collection_name(strategy: StrategyName | str, model_id: str) -> str: + """Model Corpus collection name for a Strategy + Embedding Model. + + Always includes the Embedding Model id (cloud and local). + """ if isinstance(strategy, StrategyName): name = strategy.value else: name = strategy - return f"{name}_collection" + return f"{name}__{_sanitize_model_id(model_id)}_collection" + + +def parse_collection_meta(name: str) -> dict[str, Any]: + """Infer strategy + Embedding Model from a Qdrant collection name.""" + if not name.endswith("_collection"): + return { + "strategy": None, + "embedding_model_id": None, + "is_legacy": False, + "recognized": False, + } + body = name[: -len("_collection")] + if "__" in body: + strategy, model_id = body.split("__", 1) + return { + "strategy": strategy, + "embedding_model_id": model_id, + "is_legacy": False, + "recognized": True, + } + # Pre-ADR-0021 unscoped names (no longer written) + return { + "strategy": body, + "embedding_model_id": LEGACY_CLOUD_MODEL_ID, + "is_legacy": True, + "recognized": True, + } # ── Collection management ────────────────────────────────────────── -def ensure_collection(strategy: StrategyName) -> None: - """Create the collection for a strategy if it doesn't already exist.""" +def ensure_collection( + strategy: StrategyName, + *, + model_id: str, + dimension: int, +) -> None: + """Create the Model Corpus collection for a strategy if missing.""" client = get_qdrant_client() - name = collection_name(strategy) + name = collection_name(strategy, model_id) try: existing = [c.name for c in client.get_collections().collections] if name not in existing: client.create_collection( collection_name=name, vectors_config=VectorParams( - size=VECTOR_DIMENSION, + size=dimension, distance=Distance.COSINE, ), ) - logger.info("Created Qdrant collection: %s", name) + logger.info( + "Created Qdrant collection: %s (dim=%d, model=%s)", + name, + dimension, + model_id, + ) except Exception as exc: raise QdrantError(f"Failed to create collection '{name}': {exc}") from exc -def ensure_all_collections() -> None: - """Create collections for all five strategies.""" +def ensure_all_collections(*, model_id: str, dimension: int) -> None: + """Create collections for all five strategies under one Embedding Model.""" for strategy in StrategyName: - ensure_collection(strategy) + ensure_collection(strategy, model_id=model_id, dimension=dimension) -def delete_collection(strategy: StrategyName) -> None: - """Delete a strategy's collection entirely.""" +def delete_collection(strategy: StrategyName, *, model_id: str) -> None: + """Delete a strategy's collection for an Embedding Model.""" client = get_qdrant_client() - name = collection_name(strategy) + name = collection_name(strategy, model_id) try: client.delete_collection(collection_name=name) logger.info("Deleted Qdrant collection: %s", name) @@ -77,10 +125,10 @@ def delete_collection(strategy: StrategyName) -> None: raise QdrantError(f"Failed to delete collection '{name}': {exc}") from exc -def list_collection_points(strategy: StrategyName) -> int: - """Return the number of points in a strategy's collection.""" +def list_collection_points(strategy: StrategyName, *, model_id: str) -> int: + """Return the number of points in a strategy's Model Corpus collection.""" client = get_qdrant_client() - name = collection_name(strategy) + name = collection_name(strategy, model_id) try: info = client.get_collection(collection_name=name) return info.points_count or 0 @@ -90,8 +138,13 @@ def list_collection_points(strategy: StrategyName) -> int: # ── Upsert ───────────────────────────────────────────────────────── -def upsert_chunks(chunks: list[Chunk], embeddings: list[list[float]]) -> int: - """Upsert chunks with their embeddings into the appropriate collection. +def upsert_chunks( + chunks: list[Chunk], + embeddings: list[list[float]], + *, + model_id: str, +) -> int: + """Upsert chunks with their embeddings into the Model Corpus collection. All chunks must share the same strategy_name (one collection per call). Returns the number of points upserted. @@ -105,7 +158,7 @@ def upsert_chunks(chunks: list[Chunk], embeddings: list[list[float]]) -> int: strategy = chunks[0].strategy_name client = get_qdrant_client() - name = collection_name(strategy) + name = collection_name(strategy, model_id) points = [] for chunk, embedding in zip(chunks, embeddings): @@ -133,13 +186,12 @@ def search( query_vector: list[float], top_k: int = 5, document_filter: str | None = None, + *, + model_id: str, ) -> list[dict[str, Any]]: - """Vector similarity search in a strategy's collection. - - Returns a list of {chunk_id, score, payload} dicts, ordered by score. - """ + """Vector similarity search in a strategy's Model Corpus collection.""" client = get_qdrant_client() - name = collection_name(strategy) + name = collection_name(strategy, model_id) query_filter = None if document_filter: @@ -167,24 +219,81 @@ def search( raise QdrantError(f"Search failed in '{name}': {exc}") from exc -# ── Delete by document ───────────────────────────────────────────── +def get_chunks_by_indices( + strategy: StrategyName, + document_name: str, + chunk_indices: list[int], + *, + model_id: str, +) -> dict[int, dict[str, Any]]: + """Fetch chunks by document_name + chunk_index. -def delete_document_chunks(strategy: StrategyName, document_name: str) -> int: - """Remove all chunks for a given document from a strategy's collection. - - Returns the number of points deleted. + Returns a map of chunk_index → {chunk_id, score, payload}. + Missing indices are omitted (caller skips edges). """ + if not chunk_indices: + return {} + client = get_qdrant_client() - name = collection_name(strategy) + name = collection_name(strategy, model_id) + unique_indices = sorted(set(chunk_indices)) + + try: + results = client.scroll( + collection_name=name, + scroll_filter=Filter( + must=[ + FieldCondition( + key="document_name", + match=MatchValue(value=document_name), + ), + FieldCondition( + key="chunk_index", + match=MatchAny(any=unique_indices), + ), + ] + ), + limit=max(len(unique_indices), 1), + with_payload=True, + with_vectors=False, + ) + found: dict[int, dict[str, Any]] = {} + for point in results[0]: + payload = point.payload or {} + idx = payload.get("chunk_index") + if idx is None: + continue + found[int(idx)] = { + "chunk_id": payload.get("chunk_id", str(point.id)), + "score": None, + "payload": payload, + } + return found + except Exception as exc: + raise QdrantError( + f"Failed to fetch chunks by index from '{name}': {exc}" + ) from exc + + +# ── Delete by document ───────────────────────────────────────────── + +def delete_document_chunks( + strategy: StrategyName, + document_name: str, + *, + model_id: str, +) -> int: + """Remove all chunks for a given document from a strategy collection.""" + client = get_qdrant_client() + name = collection_name(strategy, model_id) try: - # First find matching point IDs results = client.scroll( collection_name=name, scroll_filter=Filter( must=[FieldCondition(key="document_name", match=MatchValue(value=document_name))] ), - limit=10_000, # safety cap + limit=10_000, with_payload=False, with_vectors=False, ) @@ -196,16 +305,21 @@ def delete_document_chunks(strategy: StrategyName, document_name: str) -> int: collection_name=name, points_selector=PointIdsList(points=point_ids), ) - logger.info("Deleted %d points from %s for document '%s'", len(point_ids), name, document_name) + logger.info( + "Deleted %d points from %s for document '%s'", + len(point_ids), + name, + document_name, + ) return len(point_ids) except Exception as exc: raise QdrantError(f"Failed to delete from '{name}': {exc}") from exc -def delete_all_strategy_chunks(strategy: StrategyName) -> int: +def delete_all_strategy_chunks(strategy: StrategyName, *, model_id: str) -> int: """Delete all points in a strategy's collection (full wipe).""" client = get_qdrant_client() - name = collection_name(strategy) + name = collection_name(strategy, model_id) try: info = client.get_collection(collection_name=name) count = info.points_count or 0 diff --git a/src/storage/sqlite.py b/src/storage/sqlite.py index d43c57d..200738f 100644 --- a/src/storage/sqlite.py +++ b/src/storage/sqlite.py @@ -62,6 +62,7 @@ CREATE TABLE IF NOT EXISTS queries ( question TEXT NOT NULL, answer TEXT NOT NULL, retrieved_chunks TEXT NOT NULL DEFAULT '[]', + expansion_tree TEXT NOT NULL DEFAULT '[]', latency_breakdown TEXT NOT NULL DEFAULT '{}', token_usage TEXT NOT NULL DEFAULT '{}', created_at TEXT NOT NULL, @@ -76,17 +77,100 @@ CREATE TABLE IF NOT EXISTS experiments ( per_question TEXT NOT NULL DEFAULT '[]', aggregate_metrics TEXT NOT NULL DEFAULT '{}', strategies_used TEXT NOT NULL DEFAULT '[]', + embedding_model_id TEXT, + embedding_provider TEXT, + boundary_embedding_model_id TEXT, created_at TEXT NOT NULL, FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE ); + +CREATE TABLE IF NOT EXISTS app_settings ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL +); """ def init_db() -> None: - """Create tables if they don't exist.""" + """Create tables if they don't exist and apply light migrations.""" conn = _connect() try: conn.executescript(_SCHEMA) + _migrate_schema(conn) + conn.commit() + finally: + conn.close() + + +def _migrate_schema(conn: sqlite3.Connection) -> None: + """Add Embedding Model provenance columns and backfill legacy experiments.""" + cols = { + row[1] + for row in conn.execute("PRAGMA table_info(experiments)").fetchall() + } + if "embedding_model_id" not in cols: + conn.execute("ALTER TABLE experiments ADD COLUMN embedding_model_id TEXT") + if "embedding_provider" not in cols: + conn.execute("ALTER TABLE experiments ADD COLUMN embedding_provider TEXT") + if "boundary_embedding_model_id" not in cols: + conn.execute( + "ALTER TABLE experiments ADD COLUMN boundary_embedding_model_id TEXT" + ) + + dcols = { + row[1] + for row in conn.execute("PRAGMA table_info(documents)").fetchall() + } + if "last_corpus_embedding_model_id" not in dcols: + conn.execute( + "ALTER TABLE documents ADD COLUMN last_corpus_embedding_model_id TEXT" + ) + if "last_boundary_embedding_model_id" not in dcols: + conn.execute( + "ALTER TABLE documents ADD COLUMN last_boundary_embedding_model_id TEXT" + ) + + qcols = { + row[1] + for row in conn.execute("PRAGMA table_info(queries)").fetchall() + } + if "expansion_tree" not in qcols: + conn.execute( + "ALTER TABLE queries ADD COLUMN expansion_tree TEXT NOT NULL DEFAULT '[]'" + ) + + # Legacy Experiments without provenance → historical OpenAI small (ADR-0019) + from src.chunking.embedding_models import LEGACY_CLOUD_MODEL_ID, Provider + + conn.execute( + """UPDATE experiments + SET embedding_model_id = ?, embedding_provider = ? + WHERE embedding_model_id IS NULL OR embedding_model_id = ''""", + (LEGACY_CLOUD_MODEL_ID, Provider.CLOUD.value), + ) + + +def get_app_setting(key: str) -> str | None: + """Read a persisted app setting value.""" + conn = _connect() + try: + row = conn.execute( + "SELECT value FROM app_settings WHERE key = ?", (key,) + ).fetchone() + return row["value"] if row else None + finally: + conn.close() + + +def set_app_setting(key: str, value: str) -> None: + """Upsert a persisted app setting.""" + conn = _connect() + try: + conn.execute( + """INSERT INTO app_settings (key, value) VALUES (?, ?) + ON CONFLICT(key) DO UPDATE SET value = excluded.value""", + (key, value), + ) conn.commit() finally: conn.close() @@ -157,6 +241,27 @@ def update_chunk_counts(doc_id: str, chunk_counts: dict[str, int]) -> None: conn.close() +def update_process_embedding_provenance( + doc_id: str, + *, + corpus_model_id: str, + boundary_model_id: str | None = None, +) -> None: + """Record Corpus/Boundary used on the last successful process run.""" + conn = _connect() + try: + conn.execute( + """UPDATE documents + SET last_corpus_embedding_model_id = ?, + last_boundary_embedding_model_id = ? + WHERE id = ?""", + (corpus_model_id, boundary_model_id, doc_id), + ) + conn.commit() + finally: + conn.close() + + def delete_document(doc_id: str) -> bool: """Delete a document and its cascaded queries/experiments.""" conn = _connect() @@ -178,6 +283,7 @@ def save_query( question: str, answer: str, retrieved_chunks: list[dict] | None = None, + expansion_tree: list[dict] | None = None, latency_breakdown: dict[str, float] | None = None, token_usage: dict[str, int] | None = None, ) -> dict[str, Any]: @@ -188,10 +294,11 @@ def save_query( conn.execute( """INSERT INTO queries (id, document_id, strategy_name, question, answer, - retrieved_chunks, latency_breakdown, token_usage, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", + retrieved_chunks, expansion_tree, latency_breakdown, token_usage, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", (query_id, document_id, strategy_name, question, answer, json.dumps(retrieved_chunks or []), + json.dumps(expansion_tree or []), json.dumps(latency_breakdown or {}), json.dumps(token_usage or {}), _now()), @@ -250,22 +357,36 @@ def save_experiment( per_question: list[dict] | None = None, aggregate_metrics: dict | None = None, strategies_used: list[str] | None = None, + embedding_model_id: str | None = None, + embedding_provider: str | None = None, + boundary_embedding_model_id: str | None = None, ) -> dict[str, Any]: - """Insert a benchmark experiment.""" + """Insert a benchmark experiment. + + embedding_model_id is the Corpus Embedding Model (ADR-0024). + """ + from src.chunking.embedding_models import DEFAULT_CLOUD_MODEL_ID, Provider + experiment_id = experiment_id or _new_id() + model_id = embedding_model_id or DEFAULT_CLOUD_MODEL_ID + provider = embedding_provider or Provider.CLOUD.value conn = _connect() try: conn.execute( """INSERT INTO experiments (id, document_id, benchmark_config, questions, per_question, - aggregate_metrics, strategies_used, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", + aggregate_metrics, strategies_used, embedding_model_id, + embedding_provider, boundary_embedding_model_id, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", (experiment_id, document_id, json.dumps(benchmark_config or {}), json.dumps(questions or []), json.dumps(per_question or []), json.dumps(aggregate_metrics or {}), json.dumps(strategies_used or []), + model_id, + provider, + boundary_embedding_model_id, _now()), ) conn.commit() @@ -324,7 +445,7 @@ def delete_experiment(experiment_id: str) -> bool: # ── Internal helpers ─────────────────────────────────────────────── -_JSON_FIELDS = {"chunk_counts", "retrieved_chunks", "latency_breakdown", +_JSON_FIELDS = {"chunk_counts", "retrieved_chunks", "expansion_tree", "latency_breakdown", "token_usage", "document_tree", "benchmark_config", "questions", "per_question", "aggregate_metrics", "strategies_used"}