From f562b91f83108042dc9b410992eb1c41834f97ac Mon Sep 17 00:00:00 2001 From: Mahdi Bazrafshan Date: Mon, 10 Aug 2026 14:13:07 +0330 Subject: [PATCH] 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 --- src/benchmarking/benchmark_service.py | 60 +++++- src/benchmarking/evaluation.py | 9 +- src/benchmarking/models.py | 60 +++++- src/benchmarking/query_service.py | 278 +++++++++++++++++++++++--- src/benchmarking/report.py | 126 +++++++++++- src/benchmarking/routes.py | 20 ++ 6 files changed, 512 insertions(+), 41 deletions(-) diff --git a/src/benchmarking/benchmark_service.py b/src/benchmarking/benchmark_service.py index d34f208..fb8e8b0 100644 --- a/src/benchmarking/benchmark_service.py +++ b/src/benchmarking/benchmark_service.py @@ -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, diff --git a/src/benchmarking/evaluation.py b/src/benchmarking/evaluation.py index 537f831..80c993e 100644 --- a/src/benchmarking/evaluation.py +++ b/src/benchmarking/evaluation.py @@ -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) diff --git a/src/benchmarking/models.py b/src/benchmarking/models.py index 2648256..503b205 100644 --- a/src/benchmarking/models.py +++ b/src/benchmarking/models.py @@ -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") diff --git a/src/benchmarking/query_service.py b/src/benchmarking/query_service.py index dc1fcdc..54eaa7e 100644 --- a/src/benchmarking/query_service.py +++ b/src/benchmarking/query_service.py @@ -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 - context_parts.append(f"[Chunk {i} (score: {score:.3f})]") + 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"), + ) - context_parts.append(f"[Chunk {i} (score: {score:.3f})]") + 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"], diff --git a/src/benchmarking/report.py b/src/benchmarking/report.py index 92782be..c2c4800 100644 --- a/src/benchmarking/report.py +++ b/src/benchmarking/report.py @@ -175,6 +175,8 @@ def generate_managerial_report(experiment: dict) -> str:
{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}
@@ -247,6 +249,13 @@ def generate_managerial_report(experiment: dict) -> str: {_generate_decision_insights(rankings, aggregate)} + +
+

Neighbor Expansion Tree

+
+
+ {_build_expansion_tree_section(per_question, strategies, compact=True)} +

Details

@@ -306,7 +315,9 @@ def generate_technical_report(experiment: dict) -> str:
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}
@@ -444,6 +455,13 @@ def generate_technical_report(experiment: dict) -> str:
+ +
+

Neighbor Expansion Tree

+
+
+ {_build_expansion_tree_section(per_question, strategies, compact=False)} +

Navigation

@@ -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'
' + f'↑ prev {html.escape(str(n.get("chunk_index", "?")))} ' + f'{html.escape(str(n.get("text", ""))[:160])}
' + for n in (node.get("neighbors_prev") or []) + ) + next_html = "".join( + f'
' + f'↓ next {html.escape(str(n.get("chunk_index", "?")))} ' + f'{html.escape(str(n.get("text", ""))[:160])}
' + 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'
' + f'
' + f'Hit #{i + 1}' + f'{score_s}
' + f'{prev_html}' + f'
' + f'● hit {html.escape(str(node.get("chunk_index", "?")))} ' + f'{html.escape(str(node.get("text", ""))[:220])}
' + f'{next_html}
' + ) + blocks.append( + f'
' + f'
' + f'{q_id} · {html.escape(str(strategy))} · {q_text}
' + f'{"".join(hit_blocks)}
' + ) + + if not compact: + continue + # compact: stop after first question that contributed blocks + if blocks: + break + + if not blocks: + return ( + '

' + "No Expansion Tree data stored for this Experiment " + "(Neighbor Expansion off, all questions failed before a tree was saved, " + "or no fixed_size results).

" + ) + + note = "" + if compact and len(per_question) > 1: + if first_tree_index and first_tree_index > 0: + note = ( + '

' + 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.

" + ) + else: + note = ( + '

' + "Showing first question with Expansion Tree data. " + "Open technical view for all questions.

" + ) + return note + "".join(blocks) + + def _pill(score: float) -> str: """Create a score pill.""" if score >= 8: diff --git a/src/benchmarking/routes.py b/src/benchmarking/routes.py index 9da9f7f..13b2fd3 100644 --- a/src/benchmarking/routes.py +++ b/src/benchmarking/routes.py @@ -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