mba/chunking_strategies_evaluation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(storage): add app settings and model-scoped collections

Browse Source 
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 <cursoragent@cursor.com>
This commit is contained in:
Mahdi Bazrafshan
2026-08-10 14:12:45 +03:30
parent 5fd19c12d9
commit 0ff9d8dd21
2 changed files with 284 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

198
src/storage/qdrant.py
View File

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

135
src/storage/sqlite.py
View File

@@ -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"}