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 <cursoragent@cursor.com>
This commit is contained in:
@@ -1,75 +1,123 @@
|
|||||||
"""Qdrant vector storage layer.
|
"""Qdrant vector storage layer.
|
||||||
|
|
||||||
One collection per chunking strategy. Handles collection creation,
|
One collection per (Strategy, Embedding Model) — the Model Corpus.
|
||||||
vector upsert, and similarity search.
|
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 logging
|
||||||
import uuid
|
import uuid
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from qdrant_client import QdrantClient
|
|
||||||
from qdrant_client.models import (
|
from qdrant_client.models import (
|
||||||
Distance,
|
Distance,
|
||||||
FieldCondition,
|
FieldCondition,
|
||||||
Filter,
|
Filter,
|
||||||
|
MatchAny,
|
||||||
MatchValue,
|
MatchValue,
|
||||||
PointIdsList,
|
PointIdsList,
|
||||||
PointStruct,
|
PointStruct,
|
||||||
VectorParams,
|
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.dependencies import get_qdrant_client
|
||||||
from src.core.exceptions import QdrantError
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
# Embedding dimension for text-embedding-3-small
|
# Backward-compat alias — prefer EmbeddingModelSpec.dimension
|
||||||
VECTOR_DIMENSION = 1536
|
VECTOR_DIMENSION = 1536
|
||||||
|
|
||||||
|
|
||||||
def collection_name(strategy: StrategyName | str) -> str:
|
def _sanitize_model_id(model_id: str) -> str:
|
||||||
"""Convention: {strategy_name}_collection."""
|
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):
|
if isinstance(strategy, StrategyName):
|
||||||
name = strategy.value
|
name = strategy.value
|
||||||
else:
|
else:
|
||||||
name = strategy
|
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 ──────────────────────────────────────────
|
# ── Collection management ──────────────────────────────────────────
|
||||||
|
|
||||||
def ensure_collection(strategy: StrategyName) -> None:
|
def ensure_collection(
|
||||||
"""Create the collection for a strategy if it doesn't already exist."""
|
strategy: StrategyName,
|
||||||
|
*,
|
||||||
|
model_id: str,
|
||||||
|
dimension: int,
|
||||||
|
) -> None:
|
||||||
|
"""Create the Model Corpus collection for a strategy if missing."""
|
||||||
client = get_qdrant_client()
|
client = get_qdrant_client()
|
||||||
name = collection_name(strategy)
|
name = collection_name(strategy, model_id)
|
||||||
try:
|
try:
|
||||||
existing = [c.name for c in client.get_collections().collections]
|
existing = [c.name for c in client.get_collections().collections]
|
||||||
if name not in existing:
|
if name not in existing:
|
||||||
client.create_collection(
|
client.create_collection(
|
||||||
collection_name=name,
|
collection_name=name,
|
||||||
vectors_config=VectorParams(
|
vectors_config=VectorParams(
|
||||||
size=VECTOR_DIMENSION,
|
size=dimension,
|
||||||
distance=Distance.COSINE,
|
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:
|
except Exception as exc:
|
||||||
raise QdrantError(f"Failed to create collection '{name}': {exc}") from exc
|
raise QdrantError(f"Failed to create collection '{name}': {exc}") from exc
|
||||||
|
|
||||||
|
|
||||||
def ensure_all_collections() -> None:
|
def ensure_all_collections(*, model_id: str, dimension: int) -> None:
|
||||||
"""Create collections for all five strategies."""
|
"""Create collections for all five strategies under one Embedding Model."""
|
||||||
for strategy in StrategyName:
|
for strategy in StrategyName:
|
||||||
ensure_collection(strategy)
|
ensure_collection(strategy, model_id=model_id, dimension=dimension)
|
||||||
|
|
||||||
|
|
||||||
def delete_collection(strategy: StrategyName) -> None:
|
def delete_collection(strategy: StrategyName, *, model_id: str) -> None:
|
||||||
"""Delete a strategy's collection entirely."""
|
"""Delete a strategy's collection for an Embedding Model."""
|
||||||
client = get_qdrant_client()
|
client = get_qdrant_client()
|
||||||
name = collection_name(strategy)
|
name = collection_name(strategy, model_id)
|
||||||
try:
|
try:
|
||||||
client.delete_collection(collection_name=name)
|
client.delete_collection(collection_name=name)
|
||||||
logger.info("Deleted Qdrant collection: %s", 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
|
raise QdrantError(f"Failed to delete collection '{name}': {exc}") from exc
|
||||||
|
|
||||||
|
|
||||||
def list_collection_points(strategy: StrategyName) -> int:
|
def list_collection_points(strategy: StrategyName, *, model_id: str) -> int:
|
||||||
"""Return the number of points in a strategy's collection."""
|
"""Return the number of points in a strategy's Model Corpus collection."""
|
||||||
client = get_qdrant_client()
|
client = get_qdrant_client()
|
||||||
name = collection_name(strategy)
|
name = collection_name(strategy, model_id)
|
||||||
try:
|
try:
|
||||||
info = client.get_collection(collection_name=name)
|
info = client.get_collection(collection_name=name)
|
||||||
return info.points_count or 0
|
return info.points_count or 0
|
||||||
@@ -90,8 +138,13 @@ def list_collection_points(strategy: StrategyName) -> int:
|
|||||||
|
|
||||||
# ── Upsert ─────────────────────────────────────────────────────────
|
# ── Upsert ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
def upsert_chunks(chunks: list[Chunk], embeddings: list[list[float]]) -> int:
|
def upsert_chunks(
|
||||||
"""Upsert chunks with their embeddings into the appropriate collection.
|
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).
|
All chunks must share the same strategy_name (one collection per call).
|
||||||
Returns the number of points upserted.
|
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
|
strategy = chunks[0].strategy_name
|
||||||
client = get_qdrant_client()
|
client = get_qdrant_client()
|
||||||
name = collection_name(strategy)
|
name = collection_name(strategy, model_id)
|
||||||
|
|
||||||
points = []
|
points = []
|
||||||
for chunk, embedding in zip(chunks, embeddings):
|
for chunk, embedding in zip(chunks, embeddings):
|
||||||
@@ -133,13 +186,12 @@ def search(
|
|||||||
query_vector: list[float],
|
query_vector: list[float],
|
||||||
top_k: int = 5,
|
top_k: int = 5,
|
||||||
document_filter: str | None = None,
|
document_filter: str | None = None,
|
||||||
|
*,
|
||||||
|
model_id: str,
|
||||||
) -> list[dict[str, Any]]:
|
) -> list[dict[str, Any]]:
|
||||||
"""Vector similarity search in a strategy's collection.
|
"""Vector similarity search in a strategy's Model Corpus collection."""
|
||||||
|
|
||||||
Returns a list of {chunk_id, score, payload} dicts, ordered by score.
|
|
||||||
"""
|
|
||||||
client = get_qdrant_client()
|
client = get_qdrant_client()
|
||||||
name = collection_name(strategy)
|
name = collection_name(strategy, model_id)
|
||||||
|
|
||||||
query_filter = None
|
query_filter = None
|
||||||
if document_filter:
|
if document_filter:
|
||||||
@@ -167,24 +219,81 @@ def search(
|
|||||||
raise QdrantError(f"Search failed in '{name}': {exc}") from exc
|
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:
|
Returns a map of chunk_index → {chunk_id, score, payload}.
|
||||||
"""Remove all chunks for a given document from a strategy's collection.
|
Missing indices are omitted (caller skips edges).
|
||||||
|
|
||||||
Returns the number of points deleted.
|
|
||||||
"""
|
"""
|
||||||
|
if not chunk_indices:
|
||||||
|
return {}
|
||||||
|
|
||||||
client = get_qdrant_client()
|
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:
|
try:
|
||||||
# First find matching point IDs
|
|
||||||
results = client.scroll(
|
results = client.scroll(
|
||||||
collection_name=name,
|
collection_name=name,
|
||||||
scroll_filter=Filter(
|
scroll_filter=Filter(
|
||||||
must=[FieldCondition(key="document_name", match=MatchValue(value=document_name))]
|
must=[FieldCondition(key="document_name", match=MatchValue(value=document_name))]
|
||||||
),
|
),
|
||||||
limit=10_000, # safety cap
|
limit=10_000,
|
||||||
with_payload=False,
|
with_payload=False,
|
||||||
with_vectors=False,
|
with_vectors=False,
|
||||||
)
|
)
|
||||||
@@ -196,16 +305,21 @@ def delete_document_chunks(strategy: StrategyName, document_name: str) -> int:
|
|||||||
collection_name=name,
|
collection_name=name,
|
||||||
points_selector=PointIdsList(points=point_ids),
|
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)
|
return len(point_ids)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
raise QdrantError(f"Failed to delete from '{name}': {exc}") from 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)."""
|
"""Delete all points in a strategy's collection (full wipe)."""
|
||||||
client = get_qdrant_client()
|
client = get_qdrant_client()
|
||||||
name = collection_name(strategy)
|
name = collection_name(strategy, model_id)
|
||||||
try:
|
try:
|
||||||
info = client.get_collection(collection_name=name)
|
info = client.get_collection(collection_name=name)
|
||||||
count = info.points_count or 0
|
count = info.points_count or 0
|
||||||
|
|||||||
@@ -62,6 +62,7 @@ CREATE TABLE IF NOT EXISTS queries (
|
|||||||
question TEXT NOT NULL,
|
question TEXT NOT NULL,
|
||||||
answer TEXT NOT NULL,
|
answer TEXT NOT NULL,
|
||||||
retrieved_chunks TEXT NOT NULL DEFAULT '[]',
|
retrieved_chunks TEXT NOT NULL DEFAULT '[]',
|
||||||
|
expansion_tree TEXT NOT NULL DEFAULT '[]',
|
||||||
latency_breakdown TEXT NOT NULL DEFAULT '{}',
|
latency_breakdown TEXT NOT NULL DEFAULT '{}',
|
||||||
token_usage TEXT NOT NULL DEFAULT '{}',
|
token_usage TEXT NOT NULL DEFAULT '{}',
|
||||||
created_at TEXT NOT NULL,
|
created_at TEXT NOT NULL,
|
||||||
@@ -76,17 +77,100 @@ CREATE TABLE IF NOT EXISTS experiments (
|
|||||||
per_question TEXT NOT NULL DEFAULT '[]',
|
per_question TEXT NOT NULL DEFAULT '[]',
|
||||||
aggregate_metrics TEXT NOT NULL DEFAULT '{}',
|
aggregate_metrics TEXT NOT NULL DEFAULT '{}',
|
||||||
strategies_used 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,
|
created_at TEXT NOT NULL,
|
||||||
FOREIGN KEY (document_id) REFERENCES documents(id) ON DELETE CASCADE
|
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:
|
def init_db() -> None:
|
||||||
"""Create tables if they don't exist."""
|
"""Create tables if they don't exist and apply light migrations."""
|
||||||
conn = _connect()
|
conn = _connect()
|
||||||
try:
|
try:
|
||||||
conn.executescript(_SCHEMA)
|
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()
|
conn.commit()
|
||||||
finally:
|
finally:
|
||||||
conn.close()
|
conn.close()
|
||||||
@@ -157,6 +241,27 @@ def update_chunk_counts(doc_id: str, chunk_counts: dict[str, int]) -> None:
|
|||||||
conn.close()
|
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:
|
def delete_document(doc_id: str) -> bool:
|
||||||
"""Delete a document and its cascaded queries/experiments."""
|
"""Delete a document and its cascaded queries/experiments."""
|
||||||
conn = _connect()
|
conn = _connect()
|
||||||
@@ -178,6 +283,7 @@ def save_query(
|
|||||||
question: str,
|
question: str,
|
||||||
answer: str,
|
answer: str,
|
||||||
retrieved_chunks: list[dict] | None = None,
|
retrieved_chunks: list[dict] | None = None,
|
||||||
|
expansion_tree: list[dict] | None = None,
|
||||||
latency_breakdown: dict[str, float] | None = None,
|
latency_breakdown: dict[str, float] | None = None,
|
||||||
token_usage: dict[str, int] | None = None,
|
token_usage: dict[str, int] | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
@@ -188,10 +294,11 @@ def save_query(
|
|||||||
conn.execute(
|
conn.execute(
|
||||||
"""INSERT INTO queries
|
"""INSERT INTO queries
|
||||||
(id, document_id, strategy_name, question, answer,
|
(id, document_id, strategy_name, question, answer,
|
||||||
retrieved_chunks, latency_breakdown, token_usage, created_at)
|
retrieved_chunks, expansion_tree, latency_breakdown, token_usage, created_at)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||||
(query_id, document_id, strategy_name, question, answer,
|
(query_id, document_id, strategy_name, question, answer,
|
||||||
json.dumps(retrieved_chunks or []),
|
json.dumps(retrieved_chunks or []),
|
||||||
|
json.dumps(expansion_tree or []),
|
||||||
json.dumps(latency_breakdown or {}),
|
json.dumps(latency_breakdown or {}),
|
||||||
json.dumps(token_usage or {}),
|
json.dumps(token_usage or {}),
|
||||||
_now()),
|
_now()),
|
||||||
@@ -250,22 +357,36 @@ def save_experiment(
|
|||||||
per_question: list[dict] | None = None,
|
per_question: list[dict] | None = None,
|
||||||
aggregate_metrics: dict | None = None,
|
aggregate_metrics: dict | None = None,
|
||||||
strategies_used: list[str] | 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]:
|
) -> 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()
|
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()
|
conn = _connect()
|
||||||
try:
|
try:
|
||||||
conn.execute(
|
conn.execute(
|
||||||
"""INSERT INTO experiments
|
"""INSERT INTO experiments
|
||||||
(id, document_id, benchmark_config, questions, per_question,
|
(id, document_id, benchmark_config, questions, per_question,
|
||||||
aggregate_metrics, strategies_used, created_at)
|
aggregate_metrics, strategies_used, embedding_model_id,
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
embedding_provider, boundary_embedding_model_id, created_at)
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||||
(experiment_id, document_id,
|
(experiment_id, document_id,
|
||||||
json.dumps(benchmark_config or {}),
|
json.dumps(benchmark_config or {}),
|
||||||
json.dumps(questions or []),
|
json.dumps(questions or []),
|
||||||
json.dumps(per_question or []),
|
json.dumps(per_question or []),
|
||||||
json.dumps(aggregate_metrics or {}),
|
json.dumps(aggregate_metrics or {}),
|
||||||
json.dumps(strategies_used or []),
|
json.dumps(strategies_used or []),
|
||||||
|
model_id,
|
||||||
|
provider,
|
||||||
|
boundary_embedding_model_id,
|
||||||
_now()),
|
_now()),
|
||||||
)
|
)
|
||||||
conn.commit()
|
conn.commit()
|
||||||
@@ -324,7 +445,7 @@ def delete_experiment(experiment_id: str) -> bool:
|
|||||||
|
|
||||||
# ── Internal helpers ───────────────────────────────────────────────
|
# ── 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",
|
"token_usage", "document_tree", "benchmark_config",
|
||||||
"questions", "per_question", "aggregate_metrics",
|
"questions", "per_question", "aggregate_metrics",
|
||||||
"strategies_used"}
|
"strategies_used"}
|
||||||
|
|||||||
Reference in New Issue
Block a user