feat(benchmarking): add model-scoped query and neighbor expansion
Why: - Queries and Experiments must hit the Corpus Embedding Model's collections and optionally widen fixed_size context. Changes: - Resolve corpus model per request; apply Neighbor Expansion with Expansion Tree; persist and report expansion provenance. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -14,6 +14,7 @@ from typing import Any
|
||||
|
||||
from src.benchmarking.evaluation import evaluate_single
|
||||
from src.benchmarking.query_service import run_query
|
||||
from src.chunking.embedding_models import Provider
|
||||
from src.core.config import settings
|
||||
from src.core.exceptions import BenchmarkError
|
||||
from src.core.models import StrategyName
|
||||
@@ -87,15 +88,16 @@ def estimate_cost(
|
||||
) -> dict[str, Any]:
|
||||
"""Estimate the cost of running a benchmark.
|
||||
|
||||
Args:
|
||||
num_questions: Number of questions
|
||||
num_strategies: Number of strategies
|
||||
|
||||
Returns:
|
||||
Cost estimate dict
|
||||
Local Embedding Models contribute $0 embedding cost (ADR-0019).
|
||||
"""
|
||||
from src.chunking.embedding import get_corpus_embedding_model
|
||||
|
||||
active = get_corpus_embedding_model()
|
||||
|
||||
# Rough estimates based on GPT-4o-mini pricing
|
||||
embedding_cost_per_call = 0.0001
|
||||
embedding_cost_per_call = (
|
||||
0.0 if active.provider == Provider.LOCAL else 0.0001
|
||||
)
|
||||
query_cost_per_call = 0.001
|
||||
evaluation_cost_per_call = 0.001
|
||||
|
||||
@@ -118,6 +120,8 @@ def estimate_cost(
|
||||
"num_strategies": num_strategies,
|
||||
"total_queries": total_queries,
|
||||
"total_evaluations": total_evaluations,
|
||||
"embedding_model_id": active.id,
|
||||
"embedding_provider": active.provider.value,
|
||||
"estimated_tokens": {
|
||||
"input": total_input_tokens,
|
||||
"output": total_output_tokens,
|
||||
@@ -139,6 +143,9 @@ def run_benchmark(
|
||||
strategies: list[StrategyName],
|
||||
questions: list[dict],
|
||||
top_k: int = 5,
|
||||
neighbor_prev: int = 0,
|
||||
neighbor_next: int = 0,
|
||||
corpus_model_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Run a full benchmark across questions and strategies.
|
||||
|
||||
@@ -147,16 +154,40 @@ def run_benchmark(
|
||||
strategies: List of strategies to test
|
||||
questions: List of question dicts
|
||||
top_k: Number of chunks to retrieve per query
|
||||
neighbor_prev: Neighbor Expansion prev count (fixed_size only)
|
||||
neighbor_next: Neighbor Expansion next count (fixed_size only)
|
||||
corpus_model_id: Corpus Embedding Model (default Admin Corpus)
|
||||
|
||||
Returns:
|
||||
Complete benchmark results
|
||||
"""
|
||||
from src.chunking.embedding import resolve_corpus_model
|
||||
|
||||
t_start = time.time()
|
||||
embedding_model = resolve_corpus_model(corpus_model_id)
|
||||
doc = db.get_document(document_id)
|
||||
boundary_id = None
|
||||
if doc and any(
|
||||
s.value in ("semantic", "semantic_parent_child") for s in strategies
|
||||
):
|
||||
boundary_id = doc.get("last_boundary_embedding_model_id")
|
||||
|
||||
logger.info("=" * 80)
|
||||
logger.info("[BENCHMARK] Starting benchmark")
|
||||
logger.info("[BENCHMARK] Document: %s", document_id)
|
||||
logger.info("[BENCHMARK] Strategies: %s", [s.value for s in strategies])
|
||||
logger.info("[BENCHMARK] Questions: %d", len(questions))
|
||||
logger.info(
|
||||
"[BENCHMARK] Corpus Embedding Model: %s (%s)",
|
||||
embedding_model.id,
|
||||
embedding_model.provider.value,
|
||||
)
|
||||
logger.info("[BENCHMARK] Boundary (from last process): %s", boundary_id or "—")
|
||||
logger.info(
|
||||
"[BENCHMARK] Neighbor Expansion: prev=%d next=%d",
|
||||
neighbor_prev,
|
||||
neighbor_next,
|
||||
)
|
||||
|
||||
per_question_results = []
|
||||
total_cost = 0.0
|
||||
@@ -186,6 +217,9 @@ def run_benchmark(
|
||||
strategy_name=strategy,
|
||||
question=question_text,
|
||||
top_k=top_k,
|
||||
neighbor_prev=neighbor_prev,
|
||||
neighbor_next=neighbor_next,
|
||||
embedding_model=embedding_model,
|
||||
)
|
||||
t_query = time.time() - t0
|
||||
|
||||
@@ -206,6 +240,7 @@ def run_benchmark(
|
||||
question_results["strategies"][strategy.value] = {
|
||||
"answer": query_result["answer"],
|
||||
"retrieved_chunks": query_result["retrieved_chunks"],
|
||||
"expansion_tree": query_result.get("expansion_tree") or [],
|
||||
"scores": eval_scores,
|
||||
"latency": {
|
||||
"query_seconds": round(t_query, 3),
|
||||
@@ -252,11 +287,20 @@ def run_benchmark(
|
||||
"strategies": [s.value for s in strategies],
|
||||
"num_questions": len(questions),
|
||||
"top_k": top_k,
|
||||
"neighbor_prev": neighbor_prev,
|
||||
"neighbor_next": neighbor_next,
|
||||
"embedding_model_id": embedding_model.id,
|
||||
"corpus_embedding_model_id": embedding_model.id,
|
||||
"embedding_provider": embedding_model.provider.value,
|
||||
"boundary_embedding_model_id": boundary_id,
|
||||
},
|
||||
questions=questions,
|
||||
per_question=per_question_results,
|
||||
aggregate_metrics=aggregate,
|
||||
strategies_used=[s.value for s in strategies],
|
||||
embedding_model_id=embedding_model.id,
|
||||
embedding_provider=embedding_model.provider.value,
|
||||
boundary_embedding_model_id=boundary_id,
|
||||
)
|
||||
|
||||
logger.info("[BENCHMARK] Completed in %.1fs", t_total)
|
||||
@@ -267,6 +311,8 @@ def run_benchmark(
|
||||
"experiment_id": experiment["id"],
|
||||
"document_id": document_id,
|
||||
"strategies_used": [s.value for s in strategies],
|
||||
"embedding_model_id": embedding_model.id,
|
||||
"embedding_provider": embedding_model.provider.value,
|
||||
"questions_count": len(questions),
|
||||
"aggregate_metrics": aggregate,
|
||||
"best_strategy": best_strategy,
|
||||
|
||||
@@ -66,9 +66,14 @@ def _build_context_for_evaluation(retrieved_chunks: list[dict]) -> str:
|
||||
"""Build a readable context string from retrieved chunks."""
|
||||
parts = []
|
||||
for i, chunk in enumerate(retrieved_chunks, 1):
|
||||
score = chunk.get("score", 0)
|
||||
score = chunk.get("score")
|
||||
text = chunk.get("text", "")
|
||||
parts.append(f"[Chunk {i} (score: {score:.3f})]\n{text}")
|
||||
role = chunk.get("role") or ("neighbor" if score is None else "hit")
|
||||
if isinstance(score, (int, float)):
|
||||
header = f"[Chunk {i} (score: {score:.3f}, {role})]"
|
||||
else:
|
||||
header = f"[Chunk {i} (score: —, {role})]"
|
||||
parts.append(f"{header}\n{text}")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from src.core.config import settings
|
||||
from src.core.models import StrategyName
|
||||
|
||||
|
||||
@@ -15,16 +16,51 @@ class QueryRequest(BaseModel):
|
||||
strategy: StrategyName = Field(description="Chunking strategy to use")
|
||||
question: str = Field(description="Question to ask", min_length=1)
|
||||
top_k: int = Field(default=5, description="Number of chunks to retrieve", ge=1, le=20)
|
||||
neighbor_prev: int = Field(
|
||||
default_factory=lambda: settings.neighbor_prev,
|
||||
description="Neighbor Expansion: previous chunks per hit (fixed_size only)",
|
||||
ge=0,
|
||||
le=5,
|
||||
)
|
||||
neighbor_next: int = Field(
|
||||
default_factory=lambda: settings.neighbor_next,
|
||||
description="Neighbor Expansion: next chunks per hit (fixed_size only)",
|
||||
ge=0,
|
||||
le=5,
|
||||
)
|
||||
corpus_model_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Corpus Embedding Model id (query + search); default = Admin Corpus",
|
||||
)
|
||||
|
||||
|
||||
# ── Response ───────────────────────────────────────────────────────
|
||||
|
||||
class RetrievedChunk(BaseModel):
|
||||
"""A single retrieved chunk with its similarity score."""
|
||||
"""A single chunk in the flat LLM/eval context (ADR-0023)."""
|
||||
chunk_id: str
|
||||
score: Optional[float] = None
|
||||
text: str
|
||||
parent_id: Optional[str] = None
|
||||
chunk_index: Optional[int] = None
|
||||
role: Optional[str] = None # "hit" | "neighbor"
|
||||
|
||||
|
||||
class ExpansionNeighbor(BaseModel):
|
||||
"""A neighbor chunk in the Expansion Tree."""
|
||||
chunk_id: str
|
||||
text: str
|
||||
chunk_index: Optional[int] = None
|
||||
|
||||
|
||||
class ExpansionTreeNode(BaseModel):
|
||||
"""One top-k hit with its per-hit Neighbor Expansion window."""
|
||||
chunk_id: str
|
||||
score: float
|
||||
text: str
|
||||
parent_id: Optional[str] = None
|
||||
chunk_index: Optional[int] = None
|
||||
neighbors_prev: list[ExpansionNeighbor] = Field(default_factory=list)
|
||||
neighbors_next: list[ExpansionNeighbor] = Field(default_factory=list)
|
||||
|
||||
|
||||
class QueryResponse(BaseModel):
|
||||
@@ -35,6 +71,9 @@ class QueryResponse(BaseModel):
|
||||
question: str
|
||||
answer: str
|
||||
retrieved_chunks: list[RetrievedChunk]
|
||||
expansion_tree: list[ExpansionTreeNode] = Field(default_factory=list)
|
||||
neighbor_prev: int = 0
|
||||
neighbor_next: int = 0
|
||||
latency_breakdown: dict[str, float]
|
||||
token_usage: dict[str, int]
|
||||
created_at: str
|
||||
@@ -48,6 +87,7 @@ class QueryDetailResponse(BaseModel):
|
||||
question: str
|
||||
answer: str
|
||||
retrieved_chunks: list[dict]
|
||||
expansion_tree: list[dict] = Field(default_factory=list)
|
||||
latency_breakdown: dict[str, float]
|
||||
token_usage: dict[str, int]
|
||||
created_at: str
|
||||
@@ -78,6 +118,22 @@ class BenchmarkRequest(BaseModel):
|
||||
description="Path to questions JSON file (relative to project root)",
|
||||
)
|
||||
top_k: int = Field(default=5, description="Number of chunks to retrieve", ge=1, le=20)
|
||||
neighbor_prev: int = Field(
|
||||
default_factory=lambda: settings.neighbor_prev,
|
||||
description="Neighbor Expansion: previous chunks per hit (fixed_size only)",
|
||||
ge=0,
|
||||
le=5,
|
||||
)
|
||||
neighbor_next: int = Field(
|
||||
default_factory=lambda: settings.neighbor_next,
|
||||
description="Neighbor Expansion: next chunks per hit (fixed_size only)",
|
||||
ge=0,
|
||||
le=5,
|
||||
)
|
||||
corpus_model_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Corpus Embedding Model id; default = Admin Corpus",
|
||||
)
|
||||
dry_run: bool = Field(default=False, description="Only return cost estimate, don't run benchmark")
|
||||
|
||||
|
||||
|
||||
@@ -16,9 +16,10 @@ from typing import Any
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from src.chunking.embedding import embed_single
|
||||
from src.chunking.embedding import embed_single, resolve_corpus_model
|
||||
from src.chunking.embedding_models import EmbeddingModelSpec
|
||||
from src.core.config import settings
|
||||
from src.core.dependencies import get_openai_client
|
||||
from src.core.dependencies import get_openai_client, get_qdrant_client
|
||||
from src.core.exceptions import QueryError
|
||||
from src.core.models import StrategyName
|
||||
from src.storage import qdrant as qdr
|
||||
@@ -41,6 +42,8 @@ def _fetch_parent_chunks(
|
||||
child_hits: list[dict],
|
||||
strategy: StrategyName,
|
||||
document_name: str,
|
||||
*,
|
||||
model_id: str,
|
||||
) -> dict[str, dict]:
|
||||
"""Fetch parent chunks for child hits in semantic_parent_child strategy.
|
||||
|
||||
@@ -64,11 +67,10 @@ def _fetch_parent_chunks(
|
||||
# Search for parent chunks by their IDs
|
||||
parents = {}
|
||||
for parent_id in parent_ids:
|
||||
# Use Qdrant scroll to find the parent chunk
|
||||
from qdrant_client.models import Filter, FieldCondition, MatchValue
|
||||
|
||||
client = qdr.get_qdrant_client()
|
||||
name = qdr.collection_name(strategy)
|
||||
client = get_qdrant_client()
|
||||
name = qdr.collection_name(strategy, model_id)
|
||||
|
||||
try:
|
||||
results = client.scroll(
|
||||
@@ -103,6 +105,8 @@ def _build_context(
|
||||
hits: list[dict],
|
||||
strategy: StrategyName,
|
||||
document_name: str,
|
||||
*,
|
||||
model_id: str,
|
||||
) -> str:
|
||||
"""Build context string from retrieved chunks.
|
||||
|
||||
@@ -115,7 +119,9 @@ def _build_context(
|
||||
|
||||
if strategy == StrategyName.SEMANTIC_PARENT_CHILD:
|
||||
# Fetch parent chunks
|
||||
parents = _fetch_parent_chunks(hits, strategy, document_name)
|
||||
parents = _fetch_parent_chunks(
|
||||
hits, strategy, document_name, model_id=model_id
|
||||
)
|
||||
|
||||
for i, hit in enumerate(hits, 1):
|
||||
payload = hit.get("payload", {})
|
||||
@@ -124,7 +130,10 @@ def _build_context(
|
||||
parent_id = payload.get("parent_id")
|
||||
|
||||
# Add child chunk
|
||||
if isinstance(score, (int, float)):
|
||||
context_parts.append(f"[Chunk {i} (score: {score:.3f})]")
|
||||
else:
|
||||
context_parts.append(f"[Chunk {i} (score: —)]")
|
||||
context_parts.append(chunk_text)
|
||||
|
||||
# Add parent context if available
|
||||
@@ -143,10 +152,18 @@ def _build_context(
|
||||
score = hit.get("score", 0)
|
||||
chunk_text = payload.get("text", "")
|
||||
|
||||
logger.info("[CONTEXT] Chunk %d: score=%.3f, text_len=%d, chunk_id=%s",
|
||||
i, score, len(chunk_text), hit.get("chunk_id", "unknown"))
|
||||
logger.info(
|
||||
"[CONTEXT] Chunk %d: score=%s, text_len=%d, chunk_id=%s",
|
||||
i,
|
||||
f"{score:.3f}" if isinstance(score, (int, float)) else "—",
|
||||
len(chunk_text),
|
||||
hit.get("chunk_id", "unknown"),
|
||||
)
|
||||
|
||||
if isinstance(score, (int, float)):
|
||||
context_parts.append(f"[Chunk {i} (score: {score:.3f})]")
|
||||
else:
|
||||
context_parts.append(f"[Chunk {i} (score: —)]")
|
||||
context_parts.append(chunk_text)
|
||||
context_parts.append("")
|
||||
|
||||
@@ -201,6 +218,170 @@ def _generate_answer(
|
||||
raise QueryError(f"Answer generation failed: {exc}") from exc
|
||||
|
||||
|
||||
def _hit_to_chunk_dict(hit: dict, *, role: str) -> dict[str, Any]:
|
||||
payload = hit.get("payload") or {}
|
||||
score = hit.get("score")
|
||||
return {
|
||||
"chunk_id": payload.get("chunk_id", hit.get("chunk_id")),
|
||||
"score": score if role == "hit" else None,
|
||||
"text": payload.get("text", ""),
|
||||
"parent_id": payload.get("parent_id"),
|
||||
"chunk_index": payload.get("chunk_index"),
|
||||
"role": role,
|
||||
}
|
||||
|
||||
|
||||
def _neighbor_brief(entry: dict[str, Any]) -> dict[str, Any]:
|
||||
payload = entry.get("payload") or {}
|
||||
return {
|
||||
"chunk_id": payload.get("chunk_id", entry.get("chunk_id")),
|
||||
"text": payload.get("text", ""),
|
||||
"chunk_index": payload.get("chunk_index"),
|
||||
}
|
||||
|
||||
|
||||
def apply_neighbor_expansion(
|
||||
hits: list[dict],
|
||||
*,
|
||||
strategy: StrategyName,
|
||||
document_name: str,
|
||||
model_id: str,
|
||||
neighbor_prev: int,
|
||||
neighbor_next: int,
|
||||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
"""Build flat retrieved_chunks + Expansion Tree (ADR-0023).
|
||||
|
||||
Expansion runs only for fixed_size when prev/next > 0. Otherwise the tree
|
||||
lists hits with empty neighbor arrays and the flat list is the top-k hits.
|
||||
"""
|
||||
expansion_tree: list[dict[str, Any]] = []
|
||||
for hit in hits:
|
||||
payload = hit.get("payload") or {}
|
||||
expansion_tree.append({
|
||||
"chunk_id": payload.get("chunk_id", hit.get("chunk_id")),
|
||||
"score": hit.get("score", 0),
|
||||
"text": payload.get("text", ""),
|
||||
"chunk_index": payload.get("chunk_index"),
|
||||
"neighbors_prev": [],
|
||||
"neighbors_next": [],
|
||||
})
|
||||
|
||||
apply = (
|
||||
strategy == StrategyName.FIXED_SIZE
|
||||
and (neighbor_prev > 0 or neighbor_next > 0)
|
||||
and hits
|
||||
)
|
||||
if not apply:
|
||||
retrieved = [_hit_to_chunk_dict(h, role="hit") for h in hits]
|
||||
return retrieved, expansion_tree
|
||||
|
||||
needed_indices: set[int] = set()
|
||||
hit_indices: list[int | None] = []
|
||||
for hit in hits:
|
||||
payload = hit.get("payload") or {}
|
||||
idx = payload.get("chunk_index")
|
||||
hit_indices.append(idx if idx is not None else None)
|
||||
if idx is None:
|
||||
continue
|
||||
idx = int(idx)
|
||||
for d in range(1, neighbor_prev + 1):
|
||||
needed_indices.add(idx - d)
|
||||
for d in range(1, neighbor_next + 1):
|
||||
needed_indices.add(idx + d)
|
||||
|
||||
# Never fetch negative indices
|
||||
needed_indices = {i for i in needed_indices if i >= 0}
|
||||
|
||||
# Hits may themselves be neighbors of other hits — include them in the lookup map
|
||||
by_index: dict[int, dict[str, Any]] = {}
|
||||
for hit in hits:
|
||||
payload = hit.get("payload") or {}
|
||||
idx = payload.get("chunk_index")
|
||||
if idx is None:
|
||||
continue
|
||||
by_index[int(idx)] = {
|
||||
"chunk_id": payload.get("chunk_id", hit.get("chunk_id")),
|
||||
"score": hit.get("score"),
|
||||
"payload": payload,
|
||||
}
|
||||
|
||||
fetch_indices = [i for i in needed_indices if i not in by_index]
|
||||
if fetch_indices:
|
||||
by_index.update(
|
||||
qdr.get_chunks_by_indices(
|
||||
strategy,
|
||||
document_name,
|
||||
fetch_indices,
|
||||
model_id=model_id,
|
||||
)
|
||||
)
|
||||
|
||||
for tree_node, hit, idx in zip(expansion_tree, hits, hit_indices):
|
||||
if idx is None:
|
||||
continue
|
||||
idx = int(idx)
|
||||
prev_list = []
|
||||
for d in range(neighbor_prev, 0, -1):
|
||||
entry = by_index.get(idx - d)
|
||||
if entry:
|
||||
prev_list.append(_neighbor_brief(entry))
|
||||
next_list = []
|
||||
for d in range(1, neighbor_next + 1):
|
||||
entry = by_index.get(idx + d)
|
||||
if entry:
|
||||
next_list.append(_neighbor_brief(entry))
|
||||
tree_node["neighbors_prev"] = prev_list
|
||||
tree_node["neighbors_next"] = next_list
|
||||
|
||||
# Flat LLM context: hits + all neighbors, dedupe (prefer hit), sort by chunk_index
|
||||
merged: dict[str, dict[str, Any]] = {}
|
||||
for hit in hits:
|
||||
chunk = _hit_to_chunk_dict(hit, role="hit")
|
||||
cid = chunk["chunk_id"]
|
||||
if cid:
|
||||
merged[cid] = chunk
|
||||
|
||||
for node in expansion_tree:
|
||||
for nbr in node["neighbors_prev"] + node["neighbors_next"]:
|
||||
cid = nbr.get("chunk_id")
|
||||
if not cid or cid in merged:
|
||||
continue
|
||||
merged[cid] = {
|
||||
"chunk_id": cid,
|
||||
"score": None,
|
||||
"text": nbr.get("text", ""),
|
||||
"parent_id": None,
|
||||
"chunk_index": nbr.get("chunk_index"),
|
||||
"role": "neighbor",
|
||||
}
|
||||
|
||||
def _sort_key(c: dict[str, Any]) -> tuple:
|
||||
idx = c.get("chunk_index")
|
||||
if isinstance(idx, int):
|
||||
return (0, idx)
|
||||
return (1, str(c.get("chunk_id", "")))
|
||||
|
||||
retrieved = sorted(merged.values(), key=_sort_key)
|
||||
return retrieved, expansion_tree
|
||||
|
||||
|
||||
def _chunks_as_context_hits(retrieved_chunks: list[dict[str, Any]]) -> list[dict]:
|
||||
"""Adapt flat retrieved_chunks into the hit shape expected by _build_context."""
|
||||
hits = []
|
||||
for c in retrieved_chunks:
|
||||
hits.append({
|
||||
"chunk_id": c.get("chunk_id"),
|
||||
"score": c.get("score") if c.get("score") is not None else 0,
|
||||
"payload": {
|
||||
"chunk_id": c.get("chunk_id"),
|
||||
"text": c.get("text", ""),
|
||||
"parent_id": c.get("parent_id"),
|
||||
"chunk_index": c.get("chunk_index"),
|
||||
},
|
||||
})
|
||||
return hits
|
||||
|
||||
|
||||
# ── Main query function ───────────────────────────────────────────
|
||||
|
||||
def run_query(
|
||||
@@ -209,27 +390,44 @@ def run_query(
|
||||
strategy_name: StrategyName,
|
||||
question: str,
|
||||
top_k: int = 5,
|
||||
neighbor_prev: int | None = None,
|
||||
neighbor_next: int | None = None,
|
||||
embedding_model: EmbeddingModelSpec | None = None,
|
||||
corpus_model_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Run a query against a document using a specific chunking strategy.
|
||||
|
||||
Pipeline: embed → search → answer → store
|
||||
Pipeline: embed → search → Neighbor Expansion (fixed_size) → answer → store
|
||||
|
||||
Args:
|
||||
document_id: The document to query against.
|
||||
strategy_name: Which chunking strategy's collection to search.
|
||||
question: The user's question.
|
||||
top_k: Number of chunks to retrieve (default 5).
|
||||
neighbor_prev: Prev chunks per hit for fixed_size (default from config).
|
||||
neighbor_next: Next chunks per hit for fixed_size (default from config).
|
||||
embedding_model: Explicit Corpus snapshot (e.g. Experiment);
|
||||
if omitted, resolves corpus_model_id or Admin Corpus default.
|
||||
corpus_model_id: Registry id for Corpus Embedding Model.
|
||||
|
||||
Returns:
|
||||
Query result dict with answer, chunks, and metadata.
|
||||
Query result dict with answer, chunks, expansion_tree, and metadata.
|
||||
"""
|
||||
t_start = time.time()
|
||||
if embedding_model is not None:
|
||||
model = embedding_model
|
||||
else:
|
||||
model = resolve_corpus_model(corpus_model_id)
|
||||
prev_n = settings.neighbor_prev if neighbor_prev is None else neighbor_prev
|
||||
next_n = settings.neighbor_next if neighbor_next is None else neighbor_next
|
||||
logger.info("=" * 80)
|
||||
logger.info("[QUERY] Starting query pipeline")
|
||||
logger.info("[QUERY] Document ID: %s", document_id)
|
||||
logger.info("[QUERY] Strategy: %s", strategy_name.value)
|
||||
logger.info("[QUERY] Corpus Embedding Model: %s (%s)", model.id, model.provider.value)
|
||||
logger.info("[QUERY] Question: %s", question)
|
||||
logger.info("[QUERY] Top K: %d", top_k)
|
||||
logger.info("[QUERY] Neighbor Expansion: prev=%d next=%d", prev_n, next_n)
|
||||
|
||||
# 1. Load document info
|
||||
logger.info("[STEP 1] Loading document info from SQLite...")
|
||||
@@ -246,35 +444,54 @@ def run_query(
|
||||
# 2. Embed the question
|
||||
logger.info("[STEP 2] Embedding question...")
|
||||
t0 = time.time()
|
||||
question_embedding = embed_single(question)
|
||||
question_embedding = embed_single(question, model=model, purpose="query")
|
||||
t_embed = time.time() - t0
|
||||
logger.info("[STEP 2] Question embedded in %.2fs", t_embed)
|
||||
logger.info("[STEP 2] Embedding dimension: %d", len(question_embedding))
|
||||
|
||||
# 3. Vector search in Qdrant
|
||||
logger.info("[STEP 3] Searching Qdrant collection: %s_collection", strategy_name.value)
|
||||
col = qdr.collection_name(strategy_name, model.id)
|
||||
logger.info("[STEP 3] Searching Qdrant collection: %s", col)
|
||||
t1 = time.time()
|
||||
hits = qdr.search(
|
||||
strategy=strategy_name,
|
||||
query_vector=question_embedding,
|
||||
top_k=top_k,
|
||||
document_filter=document_name,
|
||||
model_id=model.id,
|
||||
)
|
||||
t_search = time.time() - t1
|
||||
logger.info("[STEP 3] Search completed in %.2fs", t_search)
|
||||
logger.info("[STEP 3] Found %d chunks", len(hits))
|
||||
|
||||
# Log each hit
|
||||
for i, hit in enumerate(hits, 1):
|
||||
payload = hit.get("payload", {})
|
||||
logger.info("[STEP 3] Hit %d: chunk_id=%s, score=%.4f, text_len=%d",
|
||||
i, hit.get("chunk_id", "unknown"), hit.get("score", 0),
|
||||
len(payload.get("text", "")))
|
||||
|
||||
# 3b. Neighbor Expansion (fixed_size) → flat list + Expansion Tree
|
||||
retrieved_chunks, expansion_tree = apply_neighbor_expansion(
|
||||
hits,
|
||||
strategy=strategy_name,
|
||||
document_name=document_name,
|
||||
model_id=model.id,
|
||||
neighbor_prev=prev_n,
|
||||
neighbor_next=next_n,
|
||||
)
|
||||
logger.info(
|
||||
"[STEP 3b] Expansion: %d flat chunks, tree nodes=%d",
|
||||
len(retrieved_chunks),
|
||||
len(expansion_tree),
|
||||
)
|
||||
|
||||
# 4. Build context and generate answer
|
||||
logger.info("[STEP 4] Building context...")
|
||||
t2 = time.time()
|
||||
context = _build_context(hits, strategy_name, document_name)
|
||||
context_hits = _chunks_as_context_hits(retrieved_chunks)
|
||||
context = _build_context(
|
||||
context_hits, strategy_name, document_name, model_id=model.id
|
||||
)
|
||||
|
||||
logger.info("[STEP 4] Generating answer...")
|
||||
client = get_openai_client()
|
||||
@@ -284,20 +501,16 @@ def run_query(
|
||||
|
||||
t_total = time.time() - t_start
|
||||
|
||||
# 5. Prepare retrieved chunks for storage
|
||||
# 5. Log retrieved chunks
|
||||
logger.info("[STEP 5] Preparing retrieved chunks for storage...")
|
||||
retrieved_chunks = []
|
||||
for hit in hits:
|
||||
payload = hit.get("payload", {})
|
||||
chunk_data = {
|
||||
"chunk_id": payload.get("chunk_id", hit.get("chunk_id")),
|
||||
"score": hit.get("score", 0),
|
||||
"text": payload.get("text", ""),
|
||||
"parent_id": payload.get("parent_id"),
|
||||
}
|
||||
retrieved_chunks.append(chunk_data)
|
||||
logger.info("[STEP 5] Chunk: id=%s, score=%.4f, text_len=%d",
|
||||
chunk_data["chunk_id"], chunk_data["score"], len(chunk_data["text"]))
|
||||
for chunk_data in retrieved_chunks:
|
||||
logger.info(
|
||||
"[STEP 5] Chunk: id=%s, role=%s, score=%s, text_len=%d",
|
||||
chunk_data["chunk_id"],
|
||||
chunk_data.get("role"),
|
||||
chunk_data.get("score"),
|
||||
len(chunk_data["text"]),
|
||||
)
|
||||
|
||||
# 6. Store query result in SQLite
|
||||
logger.info("[STEP 6] Storing query result in SQLite...")
|
||||
@@ -314,6 +527,7 @@ def run_query(
|
||||
question=question,
|
||||
answer=answer,
|
||||
retrieved_chunks=retrieved_chunks,
|
||||
expansion_tree=expansion_tree,
|
||||
latency_breakdown=latency_breakdown,
|
||||
token_usage=token_usage,
|
||||
)
|
||||
@@ -326,9 +540,15 @@ def run_query(
|
||||
"query_id": query_record["id"],
|
||||
"document_id": document_id,
|
||||
"strategy": strategy_name.value,
|
||||
"embedding_model_id": model.id,
|
||||
"corpus_embedding_model_id": model.id,
|
||||
"embedding_provider": model.provider.value,
|
||||
"question": question,
|
||||
"answer": answer,
|
||||
"retrieved_chunks": retrieved_chunks,
|
||||
"expansion_tree": expansion_tree,
|
||||
"neighbor_prev": prev_n,
|
||||
"neighbor_next": next_n,
|
||||
"latency_breakdown": latency_breakdown,
|
||||
"token_usage": token_usage,
|
||||
"created_at": query_record["created_at"],
|
||||
|
||||
@@ -175,6 +175,8 @@ def generate_managerial_report(experiment: dict) -> str:
|
||||
<div class="meta">
|
||||
{config.get('num_questions', 0)} questions ·
|
||||
{len(strategies)} strategies ·
|
||||
top_k={config.get('top_k', 5)} ·
|
||||
neighbors={config.get('neighbor_prev', 0)}/{config.get('neighbor_next', 0)} ·
|
||||
${estimated_cost:.4f} cost{source_meta}
|
||||
</div>
|
||||
</header>
|
||||
@@ -247,6 +249,13 @@ def generate_managerial_report(experiment: dict) -> str:
|
||||
{_generate_decision_insights(rankings, aggregate)}
|
||||
</div>
|
||||
|
||||
<!-- Expansion Tree -->
|
||||
<div class="section-header">
|
||||
<h2>Neighbor Expansion Tree</h2>
|
||||
<div class="line"></div>
|
||||
</div>
|
||||
{_build_expansion_tree_section(per_question, strategies, compact=True)}
|
||||
|
||||
<!-- Quick Links -->
|
||||
<div class="section-header">
|
||||
<h2>Details</h2>
|
||||
@@ -306,7 +315,9 @@ def generate_technical_report(experiment: dict) -> str:
|
||||
<div class="meta">
|
||||
Experiment: {experiment.get('id', 'N/A')[:16]}... ·
|
||||
{total_questions} questions ·
|
||||
{len(strategies)} strategies{source_meta}
|
||||
{len(strategies)} strategies ·
|
||||
top_k={top_k} ·
|
||||
neighbors={config.get('neighbor_prev', 0)}/{config.get('neighbor_next', 0)}{source_meta}
|
||||
</div>
|
||||
</header>
|
||||
|
||||
@@ -444,6 +455,13 @@ def generate_technical_report(experiment: dict) -> str:
|
||||
</table>
|
||||
</div>
|
||||
|
||||
<!-- Expansion Tree -->
|
||||
<div class="section-header">
|
||||
<h2>Neighbor Expansion Tree</h2>
|
||||
<div class="line"></div>
|
||||
</div>
|
||||
{_build_expansion_tree_section(per_question, strategies, compact=False)}
|
||||
|
||||
<!-- Quick Links -->
|
||||
<div class="section-header">
|
||||
<h2>Navigation</h2>
|
||||
@@ -977,6 +995,112 @@ def _build_token_rows(per_question: list, strategies: list) -> str:
|
||||
return rows
|
||||
|
||||
|
||||
def _question_has_expansion_tree(qr: dict, strategies: list) -> bool:
|
||||
for strategy in strategies:
|
||||
tree = (qr.get("strategies") or {}).get(strategy, {}).get("expansion_tree") or []
|
||||
if tree:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _build_expansion_tree_section(
|
||||
per_question: list,
|
||||
strategies: list,
|
||||
*,
|
||||
compact: bool = False,
|
||||
) -> str:
|
||||
"""Render Expansion Tree HTML for fixed_size (and any strategy that has a tree).
|
||||
|
||||
Compact (managerial) mode shows the first question that has tree data,
|
||||
skipping earlier failed/empty questions.
|
||||
"""
|
||||
blocks: list[str] = []
|
||||
first_tree_index: int | None = None
|
||||
|
||||
for qi, qr in enumerate(per_question):
|
||||
if compact and blocks:
|
||||
break
|
||||
if compact and not _question_has_expansion_tree(qr, strategies):
|
||||
continue
|
||||
|
||||
q_id = html.escape(str(qr.get("question_id", "")))
|
||||
q_text = html.escape(str(qr.get("question", ""))[:120])
|
||||
for strategy in strategies:
|
||||
strat = qr.get("strategies", {}).get(strategy, {})
|
||||
tree = strat.get("expansion_tree") or []
|
||||
if not tree:
|
||||
continue
|
||||
if first_tree_index is None:
|
||||
first_tree_index = qi
|
||||
hit_blocks = []
|
||||
for i, node in enumerate(tree):
|
||||
prev_html = "".join(
|
||||
f'<div style="padding:4px 8px;margin:2px 0;border-left:3px solid var(--text-muted);'
|
||||
f'color:var(--text-muted);font-size:12px;">'
|
||||
f'<span class="pill mid">↑ prev {html.escape(str(n.get("chunk_index", "?")))}</span> '
|
||||
f'{html.escape(str(n.get("text", ""))[:160])}</div>'
|
||||
for n in (node.get("neighbors_prev") or [])
|
||||
)
|
||||
next_html = "".join(
|
||||
f'<div style="padding:4px 8px;margin:2px 0;border-left:3px solid var(--text-muted);'
|
||||
f'color:var(--text-muted);font-size:12px;">'
|
||||
f'<span class="pill mid">↓ next {html.escape(str(n.get("chunk_index", "?")))}</span> '
|
||||
f'{html.escape(str(n.get("text", ""))[:160])}</div>'
|
||||
for n in (node.get("neighbors_next") or [])
|
||||
)
|
||||
score = node.get("score")
|
||||
score_s = f"{score:.3f}" if isinstance(score, (int, float)) else "—"
|
||||
hit_blocks.append(
|
||||
f'<div style="border:1px solid var(--border);border-radius:8px;'
|
||||
f'padding:10px;margin:8px 0;background:var(--surface);">'
|
||||
f'<div style="display:flex;justify-content:space-between;margin-bottom:6px;">'
|
||||
f'<span class="pill high">Hit #{i + 1}</span>'
|
||||
f'<span class="pill mid">{score_s}</span></div>'
|
||||
f'{prev_html}'
|
||||
f'<div style="padding:6px 8px;margin:2px 0;border-left:3px solid var(--accent);font-size:13px;">'
|
||||
f'<span class="pill high">● hit {html.escape(str(node.get("chunk_index", "?")))}</span> '
|
||||
f'{html.escape(str(node.get("text", ""))[:220])}</div>'
|
||||
f'{next_html}</div>'
|
||||
)
|
||||
blocks.append(
|
||||
f'<div style="margin-bottom:16px;">'
|
||||
f'<div style="font-size:13px;color:var(--text-muted);margin-bottom:6px;">'
|
||||
f'<strong>{q_id}</strong> · {html.escape(str(strategy))} · {q_text}</div>'
|
||||
f'{"".join(hit_blocks)}</div>'
|
||||
)
|
||||
|
||||
if not compact:
|
||||
continue
|
||||
# compact: stop after first question that contributed blocks
|
||||
if blocks:
|
||||
break
|
||||
|
||||
if not blocks:
|
||||
return (
|
||||
'<p style="color:var(--text-muted);font-size:13px;">'
|
||||
"No Expansion Tree data stored for this Experiment "
|
||||
"(Neighbor Expansion off, all questions failed before a tree was saved, "
|
||||
"or no fixed_size results).</p>"
|
||||
)
|
||||
|
||||
note = ""
|
||||
if compact and len(per_question) > 1:
|
||||
if first_tree_index and first_tree_index > 0:
|
||||
note = (
|
||||
'<p style="color:var(--text-muted);font-size:12px;margin-bottom:8px;">'
|
||||
f"Showing first question with Expansion Tree data "
|
||||
f"(skipped {first_tree_index} earlier question(s) with errors or empty trees). "
|
||||
"Open technical view for all questions.</p>"
|
||||
)
|
||||
else:
|
||||
note = (
|
||||
'<p style="color:var(--text-muted);font-size:12px;margin-bottom:8px;">'
|
||||
"Showing first question with Expansion Tree data. "
|
||||
"Open technical view for all questions.</p>"
|
||||
)
|
||||
return note + "".join(blocks)
|
||||
|
||||
|
||||
def _pill(score: float) -> str:
|
||||
"""Create a score pill."""
|
||||
if score >= 8:
|
||||
|
||||
@@ -42,6 +42,9 @@ async def create_query(request: QueryRequest):
|
||||
strategy_name=request.strategy,
|
||||
question=request.question,
|
||||
top_k=request.top_k,
|
||||
neighbor_prev=request.neighbor_prev,
|
||||
neighbor_next=request.neighbor_next,
|
||||
corpus_model_id=request.corpus_model_id,
|
||||
)
|
||||
|
||||
return QueryResponse(
|
||||
@@ -53,6 +56,9 @@ async def create_query(request: QueryRequest):
|
||||
retrieved_chunks=[
|
||||
RetrievedChunk(**chunk) for chunk in result["retrieved_chunks"]
|
||||
],
|
||||
expansion_tree=result.get("expansion_tree") or [],
|
||||
neighbor_prev=result.get("neighbor_prev", 0),
|
||||
neighbor_next=result.get("neighbor_next", 0),
|
||||
latency_breakdown=result["latency_breakdown"],
|
||||
token_usage=result["token_usage"],
|
||||
created_at=result["created_at"],
|
||||
@@ -73,6 +79,7 @@ async def get_query(query_id: str):
|
||||
question=result["question"],
|
||||
answer=result["answer"],
|
||||
retrieved_chunks=result["retrieved_chunks"],
|
||||
expansion_tree=result.get("expansion_tree") or [],
|
||||
latency_breakdown=result["latency_breakdown"],
|
||||
token_usage=result["token_usage"],
|
||||
created_at=result["created_at"],
|
||||
@@ -130,6 +137,9 @@ async def create_benchmark(request: BenchmarkRequest):
|
||||
strategies=request.strategies,
|
||||
questions=questions,
|
||||
top_k=request.top_k,
|
||||
neighbor_prev=request.neighbor_prev,
|
||||
neighbor_next=request.neighbor_next,
|
||||
corpus_model_id=request.corpus_model_id,
|
||||
)
|
||||
|
||||
return BenchmarkResponse(
|
||||
@@ -183,6 +193,11 @@ async def list_experiments(document_id: str | None = None):
|
||||
for item in result.get("items", []):
|
||||
doc = db.get_document(item.get("document_id", ""))
|
||||
item["document_filename"] = doc.get("filename", "Unknown") if doc else "Deleted"
|
||||
questions = item.get("questions") or []
|
||||
item["questions_count"] = (
|
||||
item.get("benchmark_config", {}).get("num_questions")
|
||||
or (len(questions) if isinstance(questions, list) else 0)
|
||||
)
|
||||
# Calculate best_strategy from aggregate_metrics
|
||||
aggs = item.get("aggregate_metrics", {})
|
||||
best_strat, best_score = "N/A", -1
|
||||
@@ -195,6 +210,11 @@ async def list_experiments(document_id: str | None = None):
|
||||
best_score = adjusted
|
||||
best_strat = strat
|
||||
item["best_strategy"] = best_strat
|
||||
# Surface embedding on list even if only in benchmark_config
|
||||
if not item.get("embedding_model_id"):
|
||||
cfg = item.get("benchmark_config") or {}
|
||||
item["embedding_model_id"] = cfg.get("embedding_model_id")
|
||||
item["embedding_provider"] = cfg.get("embedding_provider")
|
||||
return result
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user