feat(benchmarking): add query pipeline with logging and tracing

Why:
- Need query pipeline to ask questions against chunked documents
- Need comprehensive logging for debugging

Changes:
- Query service: embed question → vector search → LLM answer
- Query routes: POST /queries, GET /queries/{id}
- Query models: QueryRequest, QueryResponse, RetrievedChunk
- Added QueryError exception class
- Added STEP 1-6 logging for full traceability

src/benchmarking/__init__.py

@@ -0,0 +1 @@
+"""Benchmarking module: query pipeline and evaluation."""

src/benchmarking/models.py

@@ -0,0 +1,53 @@
+"""Request/response schemas for the Query API."""
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from src.core.models import StrategyName
+
+
+# ── Request ────────────────────────────────────────────────────────
+
+class QueryRequest(BaseModel):
+    """Body for POST /queries."""
+    document_id: str = Field(description="Document ID to query against")
+    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)
+
+
+# ── Response ───────────────────────────────────────────────────────
+
+class RetrievedChunk(BaseModel):
+    """A single retrieved chunk with its similarity score."""
+    chunk_id: str
+    score: float
+    text: str
+    parent_id: Optional[str] = None
+
+
+class QueryResponse(BaseModel):
+    """Returned after creating a new query."""
+    query_id: str
+    document_id: str
+    strategy: str
+    question: str
+    answer: str
+    retrieved_chunks: list[RetrievedChunk]
+    latency_breakdown: dict[str, float]
+    token_usage: dict[str, int]
+    created_at: str
+
+
+class QueryDetailResponse(BaseModel):
+    """Returned when retrieving a past query."""
+    id: str
+    document_id: str
+    strategy_name: str
+    question: str
+    answer: str
+    retrieved_chunks: list[dict]
+    latency_breakdown: dict[str, float]
+    token_usage: dict[str, int]
+    created_at: str

src/benchmarking/query_service.py

@@ -0,0 +1,340 @@
+"""Query service: embed question → vector search → LLM answer.
+
+Pipeline:
+  1. Embed the user's question via OpenAI
+  2. Search Qdrant for top-k similar chunks
+  3. For semantic_parent_child: also fetch parent context
+  4. Generate answer via gpt-4o-mini with retrieved chunks
+  5. Store query + answer in SQLite
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+from openai import OpenAI
+
+from src.chunking.embedding import embed_single
+from src.core.config import settings
+from src.core.dependencies import get_openai_client
+from src.core.exceptions import QueryError
+from src.core.models import StrategyName
+from src.storage import qdrant as qdr
+from src.storage import sqlite as db
+
+logger = logging.getLogger(__name__)
+
+# ── System prompt for answer generation ────────────────────────────
+
+_ANSWER_SYSTEM_PROMPT = (
+    "You are a helpful assistant that answers questions based on the provided context. "
+    "Use ONLY the information in the context to answer. If the context doesn't contain "
+    "enough information, say so clearly. Be concise and accurate."
+)
+
+
+# ── Parent-child handling ─────────────────────────────────────────
+
+def _fetch_parent_chunks(
+    child_hits: list[dict],
+    strategy: StrategyName,
+    document_name: str,
+) -> dict[str, dict]:
+    """Fetch parent chunks for child hits in semantic_parent_child strategy.
+
+    Returns a dict mapping parent_id -> parent payload.
+    """
+    logger.info("[PARENT-CHILD] Fetching parent chunks for %d child hits", len(child_hits))
+
+    parent_ids = set()
+    for hit in child_hits:
+        payload = hit.get("payload", {})
+        parent_id = payload.get("parent_id")
+        if parent_id:
+            parent_ids.add(parent_id)
+
+    logger.info("[PARENT-CHILD] Found %d unique parent IDs: %s", len(parent_ids), parent_ids)
+
+    if not parent_ids:
+        logger.info("[PARENT-CHILD] No parent IDs found, returning empty")
+        return {}
+
+    # 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)
+
+        try:
+            results = client.scroll(
+                collection_name=name,
+                scroll_filter=Filter(
+                    must=[
+                        FieldCondition(
+                            key="chunk_id",
+                            match=MatchValue(value=parent_id)
+                        )
+                    ]
+                ),
+                limit=1,
+                with_payload=True,
+                with_vectors=False,
+            )
+
+            if results[0]:
+                parent_point = results[0][0]
+                parents[parent_id] = parent_point.payload
+                logger.info("[PARENT-CHILD] Found parent %s: %d chars", parent_id, len(parent_point.payload.get("text", "")))
+            else:
+                logger.warning("[PARENT-CHILD] Parent %s not found in Qdrant", parent_id)
+        except Exception as exc:
+            logger.error("[PARENT-CHILD] Failed to fetch parent %s: %s", parent_id, exc)
+
+    logger.info("[PARENT-CHILD] Fetched %d parent chunks total", len(parents))
+    return parents
+
+
+def _build_context(
+    hits: list[dict],
+    strategy: StrategyName,
+    document_name: str,
+) -> str:
+    """Build context string from retrieved chunks.
+
+    For semantic_parent_child: includes both child and parent context.
+    For other strategies: includes child chunks only.
+    """
+    logger.info("[CONTEXT] Building context from %d hits for strategy: %s", len(hits), strategy.value)
+
+    context_parts = []
+
+    if strategy == StrategyName.SEMANTIC_PARENT_CHILD:
+        # Fetch parent chunks
+        parents = _fetch_parent_chunks(hits, strategy, document_name)
+
+        for i, hit in enumerate(hits, 1):
+            payload = hit.get("payload", {})
+            score = hit.get("score", 0)
+            chunk_text = payload.get("text", "")
+            parent_id = payload.get("parent_id")
+
+            # Add child chunk
+            context_parts.append(f"[Chunk {i} (score: {score:.3f})]")
+            context_parts.append(chunk_text)
+
+            # Add parent context if available
+            if parent_id and parent_id in parents:
+                parent_payload = parents[parent_id]
+                parent_text = parent_payload.get("text", "")
+                if parent_text:
+                    context_parts.append(f"\n[Context from parent section]")
+                    context_parts.append(parent_text)
+
+            context_parts.append("")  # blank line between chunks
+    else:
+        # Standard strategies: just use the chunks
+        for i, hit in enumerate(hits, 1):
+            payload = hit.get("payload", {})
+            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"))
+
+            context_parts.append(f"[Chunk {i} (score: {score:.3f})]")
+            context_parts.append(chunk_text)
+            context_parts.append("")
+
+    context = "\n".join(context_parts)
+    logger.info("[CONTEXT] Total context length: %d chars", len(context))
+    return context
+
+
+# ── Answer generation ─────────────────────────────────────────────
+
+def _generate_answer(
+    client: OpenAI,
+    question: str,
+    context: str,
+) -> tuple[str, dict[str, Any]]:
+    """Generate an answer using gpt-4o-mini with retrieved context.
+
+    Returns (answer, token_usage).
+    """
+    user_prompt = (
+        f"Context:\n{context}\n\n"
+        f"Question: {question}\n\n"
+        f"Answer:"
+    )
+
+    logger.info("[ANSWER] Generating answer with %s", settings.llm_model)
+    logger.info("[ANSWER] Context length: %d chars", len(context))
+    logger.info("[ANSWER] Question: %s", question)
+
+    try:
+        response = client.chat.completions.create(
+            model=settings.llm_model,
+            temperature=0.0,
+            max_tokens=1000,
+            messages=[
+                {"role": "system", "content": _ANSWER_SYSTEM_PROMPT},
+                {"role": "user", "content": user_prompt},
+            ],
+        )
+
+        answer = response.choices[0].message.content or ""
+        usage = {
+            "prompt_tokens": response.usage.prompt_tokens if response.usage else 0,
+            "completion_tokens": response.usage.completion_tokens if response.usage else 0,
+            "total_tokens": response.usage.total_tokens if response.usage else 0,
+        }
+
+        logger.info("[ANSWER] Generated answer: %d chars, %d tokens", len(answer), usage["total_tokens"])
+        return answer.strip(), usage
+    except Exception as exc:
+        logger.error("[ANSWER] Failed to generate answer: %s", exc)
+        raise QueryError(f"Answer generation failed: {exc}") from exc
+
+
+# ── Main query function ───────────────────────────────────────────
+
+def run_query(
+    *,
+    document_id: str,
+    strategy_name: StrategyName,
+    question: str,
+    top_k: int = 5,
+) -> dict[str, Any]:
+    """Run a query against a document using a specific chunking strategy.
+
+    Pipeline: embed → search → 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).
+
+    Returns:
+        Query result dict with answer, chunks, and metadata.
+    """
+    t_start = time.time()
+    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] Question: %s", question)
+    logger.info("[QUERY] Top K: %d", top_k)
+
+    # 1. Load document info
+    logger.info("[STEP 1] Loading document info from SQLite...")
+    doc = db.get_document(document_id)
+    if doc is None:
+        logger.error("[STEP 1] Document not found: %s", document_id)
+        raise QueryError(f"Document not found: {document_id}")
+
+    document_name = doc["filename"]
+    chunk_counts = doc.get("chunk_counts", {})
+    logger.info("[STEP 1] Document found: %s", document_name)
+    logger.info("[STEP 1] Chunk counts: %s", chunk_counts)
+
+    # 2. Embed the question
+    logger.info("[STEP 2] Embedding question...")
+    t0 = time.time()
+    question_embedding = embed_single(question)
+    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)
+    t1 = time.time()
+    hits = qdr.search(
+        strategy=strategy_name,
+        query_vector=question_embedding,
+        top_k=top_k,
+        document_filter=document_name,
+    )
+    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", "")))
+
+    # 4. Build context and generate answer
+    logger.info("[STEP 4] Building context...")
+    t2 = time.time()
+    context = _build_context(hits, strategy_name, document_name)
+
+    logger.info("[STEP 4] Generating answer...")
+    client = get_openai_client()
+    answer, token_usage = _generate_answer(client, question, context)
+    t_answer = time.time() - t2
+    logger.info("[STEP 4] Answer generated in %.2fs", t_answer)
+
+    t_total = time.time() - t_start
+
+    # 5. Prepare retrieved chunks for storage
+    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"]))
+
+    # 6. Store query result in SQLite
+    logger.info("[STEP 6] Storing query result in SQLite...")
+    latency_breakdown = {
+        "embed_seconds": round(t_embed, 3),
+        "search_seconds": round(t_search, 3),
+        "answer_seconds": round(t_answer, 3),
+        "total_seconds": round(t_total, 3),
+    }
+
+    query_record = db.save_query(
+        document_id=document_id,
+        strategy_name=strategy_name.value,
+        question=question,
+        answer=answer,
+        retrieved_chunks=retrieved_chunks,
+        latency_breakdown=latency_breakdown,
+        token_usage=token_usage,
+    )
+
+    logger.info("[STEP 6] Query stored with ID: %s", query_record["id"])
+    logger.info("[QUERY] Pipeline completed in %.2fs", t_total)
+    logger.info("=" * 80)
+
+    return {
+        "query_id": query_record["id"],
+        "document_id": document_id,
+        "strategy": strategy_name.value,
+        "question": question,
+        "answer": answer,
+        "retrieved_chunks": retrieved_chunks,
+        "latency_breakdown": latency_breakdown,
+        "token_usage": token_usage,
+        "created_at": query_record["created_at"],
+    }
+
+
+def get_query(query_id: str) -> dict[str, Any] | None:
+    """Retrieve a past query by ID."""
+    return db.get_query(query_id)

src/benchmarking/routes.py

@@ -0,0 +1,68 @@
+"""Query API routes.
+
+Endpoints:
+  POST   /queries              Ask a question against a strategy
+  GET    /queries/{id}         Retrieve a past query
+"""
+
+from fastapi import APIRouter
+
+from src.core.exceptions import QueryError
+from src.core.models import StrategyName
+from src.benchmarking import query_service
+from src.benchmarking.models import (
+    QueryRequest,
+    QueryResponse,
+    QueryDetailResponse,
+    RetrievedChunk,
+)
+
+router = APIRouter()
+
+
+@router.post("/queries", response_model=QueryResponse, status_code=201)
+async def create_query(request: QueryRequest):
+    """Ask a question against a document using a specific chunking strategy.
+
+    Pipeline: embed question → vector search → generate answer → store result.
+    """
+    result = query_service.run_query(
+        document_id=request.document_id,
+        strategy_name=request.strategy,
+        question=request.question,
+        top_k=request.top_k,
+    )
+
+    return QueryResponse(
+        query_id=result["query_id"],
+        document_id=result["document_id"],
+        strategy=result["strategy"],
+        question=result["question"],
+        answer=result["answer"],
+        retrieved_chunks=[
+            RetrievedChunk(**chunk) for chunk in result["retrieved_chunks"]
+        ],
+        latency_breakdown=result["latency_breakdown"],
+        token_usage=result["token_usage"],
+        created_at=result["created_at"],
+    )
+
+
+@router.get("/queries/{query_id}", response_model=QueryDetailResponse)
+async def get_query(query_id: str):
+    """Retrieve a past query by ID."""
+    result = query_service.get_query(query_id)
+    if result is None:
+        raise QueryError(f"Query not found: {query_id}")
+
+    return QueryDetailResponse(
+        id=result["id"],
+        document_id=result["document_id"],
+        strategy_name=result["strategy_name"],
+        question=result["question"],
+        answer=result["answer"],
+        retrieved_chunks=result["retrieved_chunks"],
+        latency_breakdown=result["latency_breakdown"],
+        token_usage=result["token_usage"],
+        created_at=result["created_at"],
+    )

src/core/exceptions.py

@@ -32,6 +32,10 @@ class BenchmarkError(Exception):
     """Base exception for benchmarking-related errors."""
 
 
+class QueryError(BenchmarkError):
+    """Raised when a query operation fails."""
+
+
 class DryRunError(BenchmarkError):
     """Raised when a dry-run estimation fails."""