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:
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