"""Query endpoint for natural language Q&A.""" import asyncio import time from uuid import UUID import structlog from fastapi import APIRouter, HTTPException from pymilvus import connections from sqlalchemy import select from src.api.dependencies import ApiKeyDep, AsyncSessionDep from src.api.tenant_context import TenantContextDep, set_rls_tenant from src.api.schemas.query import QueryRequest, QueryResponse from src.core.exceptions import DocumentNotFoundError, QueryError, RetrievalError from src.db.models import Document from src.retrieval.answer_generator import AnswerGenerator from src.retrieval.semantic_retriever import SemanticRetriever logger = structlog.get_logger(__name__) router = APIRouter(prefix="/query", tags=["Query"]) async def validate_document_ids( document_ids: list[UUID], session: AsyncSessionDep, tenant_id: UUID | None = None, node_id: UUID | None = None, ) -> None: """Validate that all document IDs exist and are completed. Args: document_ids: List of document IDs to validate session: Database session tenant_id: Tenant UUID for scoped validation node_id: Node UUID for scoped validation Raises: DocumentNotFoundError: If any document doesn't exist or isn't completed """ for doc_id in document_ids: query = select(Document).where( Document.id == doc_id, Document.status == "completed", ) if tenant_id is not None: query = query.where(Document.tenant_id == tenant_id) if node_id is not None: query = query.where(Document.node_id == node_id) result = await session.execute(query) doc = result.scalar_one_or_none() if not doc: raise DocumentNotFoundError( f"Document {doc_id} not found or not processed", details={"document_id": str(doc_id)}, ) async def process_query( request: QueryRequest, session: AsyncSessionDep, tenant_id: str = "", node_id: str = "", ) -> QueryResponse: """Process a query request. Args: request: Query request with query text and optional filters session: Database session Returns: Query response with answer, sources, and metadata Raises: DocumentNotFoundError: If document_ids are invalid RetrievalError: If retrieval fails QueryError: If answer generation fails """ start_time = time.time() # Validate document IDs if provided (scoped by tenant + node) if request.document_ids: from uuid import UUID as _UUID _tid = _UUID(tenant_id) if tenant_id else None _nid = _UUID(node_id) if node_id else None await validate_document_ids(request.document_ids, session, _tid, _nid) # Retrieve relevant sources retriever = SemanticRetriever(top_k=5) # Convert tag_filters to dict if provided tag_filter_dict = None if request.tag_filters: tag_filter_dict = {} if request.tag_filters.contains_any: tag_filter_dict["contains_any"] = request.tag_filters.contains_any if request.tag_filters.contains_all: tag_filter_dict["contains_all"] = request.tag_filters.contains_all sources, top_similarity = await retriever.retrieve( query=request.query, tenant_id=tenant_id, node_id=node_id, document_ids=request.document_ids, tag_filters=tag_filter_dict, ) # Generate answer generator = AnswerGenerator() answer, confidence = await generator.generate( query=request.query, sources=sources if request.include_context else [], top_similarity=top_similarity, ) # Calculate processing time processing_time_ms = int((time.time() - start_time) * 1000) return QueryResponse( query=request.query, answer=answer, confidence=confidence, sources=sources if request.include_context else [], processing_time_ms=processing_time_ms, ) @router.post( "", status_code=200, response_model=QueryResponse, responses={ 200: { "description": "Query processed successfully", "model": QueryResponse, }, 400: { "description": "Invalid request", "content": { "application/json": { "examples": { "empty_query": { "summary": "Empty query", "value": { "detail": [ { "loc": ["body", "query"], "msg": "Query cannot be empty", "type": "value_error", } ] }, }, "query_too_long": { "summary": "Query exceeds length limit", "value": { "detail": [ { "loc": ["body", "query"], "msg": "ensure this value has at most 1000 characters", "type": "value_error.any_str.max_length", } ] }, }, } } }, }, 401: { "description": "Authentication required", }, 404: { "description": "Document not found", "content": { "application/json": { "example": { "error": { "code": "NOT_FOUND", "message": "Document {id} not found or not processed", "details": {"document_id": "uuid"}, } } } }, }, 504: { "description": "Query processing timed out", "content": { "application/json": { "example": { "error": { "code": "TIMEOUT", "message": "Query processing timed out (>5 seconds)", } } } }, }, }, ) async def query_documents( request: QueryRequest, api_key: ApiKeyDep, session: AsyncSessionDep, ctx: TenantContextDep, ) -> QueryResponse: """Process a natural language query against uploaded documents. Query documents using semantic search and return AI-generated answers with source references from the knowledge base. **Features:** - Semantic search using vector similarity - Natural language answer generation - Source references with document context - Optional document filtering - Processing time tracking **Example Request:** ```json { "query": "What does item code 2024-4_0019 appear in?", "document_ids": ["uuid1", "uuid2"], "include_context": true } ``` **Example Response:** ```json { "query": "What does item code 2024-4_0019 appear in?", "answer": "Item code 2024-4_0019 appears in the Screen column...", "confidence": 0.92, "sources": [ { "file": "distribution_spec.xlsx", "sheet": "Item Mapping", "location": "A42:D42", "context": "2024-4_0019 | - | ○ | Screen display item" } ], "processing_time_ms": 1250 } ``` """ try: logger.info( "query_received", query=request.query, document_ids=request.document_ids, include_context=request.include_context, api_key_id=str(api_key.id), ) # Set RLS session variable for PostgreSQL defense-in-depth await set_rls_tenant(session, ctx.tenant_id) # Process query with 5-second timeout try: response = await asyncio.wait_for( process_query( request, session, tenant_id=str(ctx.tenant_id), node_id=str(ctx.node_id), ), timeout=5.0, ) logger.info( "query_completed", query=request.query, processing_time_ms=response.processing_time_ms, confidence=response.confidence, sources_count=len(response.sources), ) return response except asyncio.TimeoutError: logger.warning("query_timeout", query=request.query) raise HTTPException( status_code=504, detail={ "error": { "code": "TIMEOUT", "message": "Query processing timed out (>5 seconds)", } }, ) except DocumentNotFoundError as e: logger.warning( "document_not_found", query=request.query, error=str(e), details=e.details, ) raise HTTPException(status_code=404, detail=e.to_dict()["error"]) except (RetrievalError, QueryError) as e: logger.error( "query_processing_error", query=request.query, error=str(e), error_code=e.code, ) raise HTTPException(status_code=500, detail=e.to_dict()["error"]) except Exception as e: logger.exception( "query_unexpected_error", query=request.query, error=str(e), ) # Ensure Milvus connection is cleaned up try: connections.disconnect("default") except Exception: pass raise HTTPException( status_code=500, detail={ "error": { "code": "INTERNAL_ERROR", "message": "An unexpected error occurred", } }, )