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:
2026-08-10 14:13:07 +03:30
parent 4dd3125318
commit f562b91f83
6 changed files with 512 additions and 41 deletions

View File

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