diff --git a/src/benchmarking/__init__.py b/src/benchmarking/__init__.py new file mode 100644 index 0000000..2bf65da --- /dev/null +++ b/src/benchmarking/__init__.py @@ -0,0 +1 @@ +"""Benchmarking module: query pipeline and evaluation.""" diff --git a/src/benchmarking/models.py b/src/benchmarking/models.py new file mode 100644 index 0000000..20d7bb4 --- /dev/null +++ b/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 diff --git a/src/benchmarking/query_service.py b/src/benchmarking/query_service.py new file mode 100644 index 0000000..dc1fcdc --- /dev/null +++ b/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) diff --git a/src/benchmarking/routes.py b/src/benchmarking/routes.py new file mode 100644 index 0000000..fbed4cf --- /dev/null +++ b/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"], + ) diff --git a/src/core/exceptions.py b/src/core/exceptions.py index ae944e6..9e28acc 100644 --- a/src/core/exceptions.py +++ b/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."""