From 4aaaccec49e2366e6d7842056c6607aaf6b9d3b2 Mon Sep 17 00:00:00 2001 From: Mahdi Bazrafshan Date: Sun, 26 Jul 2026 09:37:32 +0330 Subject: [PATCH] feat(storage): add SQLite and Qdrant storage layers Why: - SQLite stores structured data (documents, queries, experiments) - Qdrant stores vector embeddings for similarity search Changes: - SQLite: schema init, document/query/experiment CRUD, JSON field parsing - Qdrant: collection management, vector upsert, similarity search, delete operations --- src/storage/__init__.py | 1 + src/storage/qdrant.py | 219 ++++++++++++++++++++++++++ src/storage/sqlite.py | 331 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 551 insertions(+) create mode 100644 src/storage/__init__.py create mode 100644 src/storage/qdrant.py create mode 100644 src/storage/sqlite.py diff --git a/src/storage/__init__.py b/src/storage/__init__.py new file mode 100644 index 0000000..b81f7c5 --- /dev/null +++ b/src/storage/__init__.py @@ -0,0 +1 @@ +"""Storage layer — SQLite persistence and Qdrant vector store.""" diff --git a/src/storage/qdrant.py b/src/storage/qdrant.py new file mode 100644 index 0000000..1c066f4 --- /dev/null +++ b/src/storage/qdrant.py @@ -0,0 +1,219 @@ +"""Qdrant vector storage layer. + +One collection per chunking strategy. Handles collection creation, +vector upsert, and similarity search. +""" + +import logging +import uuid +from typing import Any + +from qdrant_client import QdrantClient +from qdrant_client.models import ( + Distance, + FieldCondition, + Filter, + MatchValue, + PointIdsList, + PointStruct, + VectorParams, +) + +from src.core.config import settings +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 + +logger = logging.getLogger(__name__) + +# Embedding dimension for text-embedding-3-small +VECTOR_DIMENSION = 1536 + + +def collection_name(strategy: StrategyName | str) -> str: + """Convention: {strategy_name}_collection.""" + if isinstance(strategy, StrategyName): + name = strategy.value + else: + name = strategy + return f"{name}_collection" + + +# ── Collection management ────────────────────────────────────────── + +def ensure_collection(strategy: StrategyName) -> None: + """Create the collection for a strategy if it doesn't already exist.""" + client = get_qdrant_client() + name = collection_name(strategy) + 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, + distance=Distance.COSINE, + ), + ) + logger.info("Created Qdrant collection: %s", name) + 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.""" + for strategy in StrategyName: + ensure_collection(strategy) + + +def delete_collection(strategy: StrategyName) -> None: + """Delete a strategy's collection entirely.""" + client = get_qdrant_client() + name = collection_name(strategy) + try: + client.delete_collection(collection_name=name) + logger.info("Deleted Qdrant collection: %s", name) + except Exception as exc: + 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.""" + client = get_qdrant_client() + name = collection_name(strategy) + try: + info = client.get_collection(collection_name=name) + return info.points_count or 0 + except Exception as exc: + raise QdrantError(f"Failed to get info for '{name}': {exc}") from exc + + +# ── Upsert ───────────────────────────────────────────────────────── + +def upsert_chunks(chunks: list[Chunk], embeddings: list[list[float]]) -> int: + """Upsert chunks with their embeddings into the appropriate collection. + + All chunks must share the same strategy_name (one collection per call). + Returns the number of points upserted. + """ + if not chunks: + return 0 + if len(chunks) != len(embeddings): + raise QdrantError( + f"Mismatch: {len(chunks)} chunks but {len(embeddings)} embeddings" + ) + + strategy = chunks[0].strategy_name + client = get_qdrant_client() + name = collection_name(strategy) + + points = [] + for chunk, embedding in zip(chunks, embeddings): + meta = chunk_to_metadata(chunk) + points.append( + PointStruct( + id=str(uuid.uuid5(uuid.NAMESPACE_URL, chunk.chunk_id)), + vector=embedding, + payload=meta.model_dump(), + ) + ) + + try: + client.upsert(collection_name=name, points=points) + logger.info("Upserted %d points into %s", len(points), name) + return len(points) + except Exception as exc: + raise QdrantError(f"Failed to upsert into '{name}': {exc}") from exc + + +# ── Search ───────────────────────────────────────────────────────── + +def search( + strategy: StrategyName, + query_vector: list[float], + top_k: int = 5, + document_filter: str | None = None, +) -> list[dict[str, Any]]: + """Vector similarity search in a strategy's collection. + + Returns a list of {chunk_id, score, payload} dicts, ordered by score. + """ + client = get_qdrant_client() + name = collection_name(strategy) + + query_filter = None + if document_filter: + query_filter = Filter( + must=[FieldCondition(key="document_name", match=MatchValue(value=document_filter))] + ) + + try: + results = client.query_points( + collection_name=name, + query=query_vector, + limit=top_k, + query_filter=query_filter, + ) + + hits = [] + for point in results.points: + hits.append({ + "chunk_id": point.id, + "score": point.score, + "payload": point.payload, + }) + return hits + except Exception as exc: + raise QdrantError(f"Search failed in '{name}': {exc}") from exc + + +# ── Delete by document ───────────────────────────────────────────── + +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. + """ + client = get_qdrant_client() + name = collection_name(strategy) + + 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 + with_payload=False, + with_vectors=False, + ) + point_ids = [p.id for p in results[0]] + if not point_ids: + return 0 + + client.delete( + 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) + 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: + """Delete all points in a strategy's collection (full wipe).""" + client = get_qdrant_client() + name = collection_name(strategy) + try: + info = client.get_collection(collection_name=name) + count = info.points_count or 0 + if count > 0: + client.delete( + collection_name=name, + points_selector=PointIdsList(points=list(range(count))), + ) + return count + except Exception as exc: + raise QdrantError(f"Failed to wipe '{name}': {exc}") from exc diff --git a/src/storage/sqlite.py b/src/storage/sqlite.py new file mode 100644 index 0000000..21ce59b --- /dev/null +++ b/src/storage/sqlite.py @@ -0,0 +1,331 @@ +"""SQLite storage layer for documents, experiments, and queries. + +All structured data (parsed documents, benchmark results, query history) +lives here. Vectors live in Qdrant — never in SQLite. +""" + +import json +import sqlite3 +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Optional + +from src.core.config import settings + + +# ── Helpers ──────────────────────────────────────────────────────── + +def _new_id() -> str: + return uuid.uuid4().hex + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +# ── Connection ───────────────────────────────────────────────────── + +def _get_db_path() -> str: + """Extract the file path from the database_url setting.""" + # settings.database_url is like "sqlite:///./data/chunking_benchmark.db" + return settings.database_url.replace("sqlite:///", "") + + +def _connect() -> sqlite3.Connection: + """Return a new SQLite connection with row_factory.""" + db_path = _get_db_path() + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA foreign_keys=ON") + return conn + + +# ── Schema init ──────────────────────────────────────────────────── + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS documents ( + id TEXT PRIMARY KEY, + filename TEXT NOT NULL, + parsed_text TEXT NOT NULL, + document_tree TEXT NOT NULL, + chunk_counts TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS queries ( + id TEXT PRIMARY KEY, + document_id TEXT NOT NULL, + strategy_name TEXT NOT NULL, + question TEXT NOT NULL, + answer TEXT NOT NULL, + retrieved_chunks TEXT NOT NULL DEFAULT '[]', + latency_breakdown TEXT NOT NULL DEFAULT '{}', + token_usage TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL, + FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS experiments ( + id TEXT PRIMARY KEY, + document_id TEXT NOT NULL, + benchmark_config TEXT NOT NULL DEFAULT '{}', + questions TEXT NOT NULL DEFAULT '[]', + per_question TEXT NOT NULL DEFAULT '[]', + aggregate_metrics TEXT NOT NULL DEFAULT '{}', + strategies_used TEXT NOT NULL DEFAULT '[]', + created_at TEXT NOT NULL, + FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE +); +""" + + +def init_db() -> None: + """Create tables if they don't exist.""" + conn = _connect() + try: + conn.executescript(_SCHEMA) + conn.commit() + finally: + conn.close() + + +# ── Document CRUD ────────────────────────────────────────────────── + +def save_document( + *, + doc_id: str | None = None, + filename: str, + parsed_text: str, + document_tree: str, + chunk_counts: dict[str, int] | None = None, +) -> dict[str, Any]: + """Insert a parsed document. Returns the full row as a dict.""" + doc_id = doc_id or _new_id() + conn = _connect() + try: + conn.execute( + """INSERT INTO documents (id, filename, parsed_text, document_tree, chunk_counts, created_at) + VALUES (?, ?, ?, ?, ?, ?)""", + (doc_id, filename, parsed_text, document_tree, + json.dumps(chunk_counts or {}), _now()), + ) + conn.commit() + return get_document(doc_id) # type: ignore[return-value] + finally: + conn.close() + + +def get_document(doc_id: str) -> dict[str, Any] | None: + """Fetch a document by ID. Parses JSON fields back to Python objects.""" + conn = _connect() + try: + row = conn.execute("SELECT * FROM documents WHERE id = ?", (doc_id,)).fetchone() + if row is None: + return None + return _row_to_dict(row) + finally: + conn.close() + + +def list_documents(*, offset: int = 0, limit: int = 50) -> dict[str, Any]: + """List documents with pagination.""" + conn = _connect() + try: + total = conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0] + rows = conn.execute( + "SELECT * FROM documents ORDER BY created_at DESC LIMIT ? OFFSET ?", + (limit, offset), + ).fetchall() + return {"items": [_row_to_dict(r) for r in rows], "total": total, "offset": offset, "limit": limit} + finally: + conn.close() + + +def update_chunk_counts(doc_id: str, chunk_counts: dict[str, int]) -> None: + """Update the per-strategy chunk counts after processing.""" + conn = _connect() + try: + conn.execute( + "UPDATE documents SET chunk_counts = ? WHERE id = ?", + (json.dumps(chunk_counts), doc_id), + ) + conn.commit() + finally: + conn.close() + + +def delete_document(doc_id: str) -> bool: + """Delete a document and its cascaded queries/experiments.""" + conn = _connect() + try: + cur = conn.execute("DELETE FROM documents WHERE id = ?", (doc_id,)) + conn.commit() + return cur.rowcount > 0 + finally: + conn.close() + + +# ── Query CRUD ───────────────────────────────────────────────────── + +def save_query( + *, + query_id: str | None = None, + document_id: str, + strategy_name: str, + question: str, + answer: str, + retrieved_chunks: list[dict] | None = None, + latency_breakdown: dict[str, float] | None = None, + token_usage: dict[str, int] | None = None, +) -> dict[str, Any]: + """Insert a query result.""" + query_id = query_id or _new_id() + conn = _connect() + try: + conn.execute( + """INSERT INTO queries + (id, document_id, strategy_name, question, answer, + retrieved_chunks, latency_breakdown, token_usage, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""", + (query_id, document_id, strategy_name, question, answer, + json.dumps(retrieved_chunks or []), + json.dumps(latency_breakdown or {}), + json.dumps(token_usage or {}), + _now()), + ) + conn.commit() + return get_query(query_id) # type: ignore[return-value] + finally: + conn.close() + + +def get_query(query_id: str) -> dict[str, Any] | None: + """Fetch a query by ID.""" + conn = _connect() + try: + row = conn.execute("SELECT * FROM queries WHERE id = ?", (query_id,)).fetchone() + if row is None: + return None + return _row_to_dict(row) + finally: + conn.close() + + +def list_queries( + *, document_id: str | None = None, offset: int = 0, limit: int = 50 +) -> dict[str, Any]: + """List queries, optionally filtered by document.""" + conn = _connect() + try: + if document_id: + total = conn.execute( + "SELECT COUNT(*) FROM queries WHERE document_id = ?", (document_id,) + ).fetchone()[0] + rows = conn.execute( + "SELECT * FROM queries WHERE document_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?", + (document_id, limit, offset), + ).fetchall() + else: + total = conn.execute("SELECT COUNT(*) FROM queries").fetchone()[0] + rows = conn.execute( + "SELECT * FROM queries ORDER BY created_at DESC LIMIT ? OFFSET ?", + (limit, offset), + ).fetchall() + return {"items": [_row_to_dict(r) for r in rows], "total": total, "offset": offset, "limit": limit} + finally: + conn.close() + + +# ── Experiment (Benchmark) CRUD ──────────────────────────────────── + +def save_experiment( + *, + experiment_id: str | None = None, + document_id: str, + benchmark_config: dict | None = None, + questions: list[dict] | None = None, + per_question: list[dict] | None = None, + aggregate_metrics: dict | None = None, + strategies_used: list[str] | None = None, +) -> dict[str, Any]: + """Insert a benchmark experiment.""" + experiment_id = experiment_id or _new_id() + conn = _connect() + try: + conn.execute( + """INSERT INTO experiments + (id, document_id, benchmark_config, questions, per_question, + aggregate_metrics, strategies_used, 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 []), + _now()), + ) + conn.commit() + return get_experiment(experiment_id) # type: ignore[return-value] + finally: + conn.close() + + +def get_experiment(experiment_id: str) -> dict[str, Any] | None: + """Fetch an experiment by ID.""" + conn = _connect() + try: + row = conn.execute("SELECT * FROM experiments WHERE id = ?", (experiment_id,)).fetchone() + if row is None: + return None + return _row_to_dict(row) + finally: + conn.close() + + +def list_experiments( + *, document_id: str | None = None, offset: int = 0, limit: int = 50 +) -> dict[str, Any]: + """List experiments, optionally filtered by document.""" + conn = _connect() + try: + if document_id: + total = conn.execute( + "SELECT COUNT(*) FROM experiments WHERE document_id = ?", (document_id,) + ).fetchone()[0] + rows = conn.execute( + "SELECT * FROM experiments WHERE document_id = ? ORDER BY created_at DESC LIMIT ? OFFSET ?", + (document_id, limit, offset), + ).fetchall() + else: + total = conn.execute("SELECT COUNT(*) FROM experiments").fetchone()[0] + rows = conn.execute( + "SELECT * FROM experiments ORDER BY created_at DESC LIMIT ? OFFSET ?", + (limit, offset), + ).fetchall() + return {"items": [_row_to_dict(r) for r in rows], "total": total, "offset": offset, "limit": limit} + finally: + conn.close() + + +# ── Internal helpers ─────────────────────────────────────────────── + +_JSON_FIELDS = {"chunk_counts", "retrieved_chunks", "latency_breakdown", + "token_usage", "document_tree", "benchmark_config", + "questions", "per_question", "aggregate_metrics", + "strategies_used"} + + +def _row_to_dict(row: sqlite3.Row) -> dict[str, Any]: + """Convert a sqlite3.Row to a plain dict, parsing JSON columns.""" + d = dict(row) + for key in _JSON_FIELDS: + if key in d and isinstance(d[key], str): + try: + d[key] = json.loads(d[key]) + except (json.JSONDecodeError, TypeError): + pass + return d