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
This commit is contained in:
1
src/storage/__init__.py
Normal file
1
src/storage/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Storage layer — SQLite persistence and Qdrant vector store."""
|
||||
219
src/storage/qdrant.py
Normal file
219
src/storage/qdrant.py
Normal file
@@ -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
|
||||
331
src/storage/sqlite.py
Normal file
331
src/storage/sqlite.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user