app/api/query.py

"""
Ukweli — RAG Query Endpoint
POST /query — The primary retrieval-augmented generation endpoint.
Orchestrates the full pipeline per Architecture Section 4.2.
"""

from __future__ import annotations

import logging
import time
import uuid

from fastapi import APIRouter, Depends
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.auth import AuthContext, require_auth
from app.db.session import get_db_session
from app.dependencies import get_citation_formatter, get_llm_gateway, get_retrieval_orchestrator
from app.models.database import QueryLog
from app.models.enums import Confidence
from app.models.schemas import (
    QueryRequest,
    QueryResponse,
    RetrievalMetadata,
)
from app.services.citation.formatter import CitationFormatter
from app.services.llm.gateway import LLMGateway
from app.services.llm.guardrails import check_query_safety
from app.services.llm.prompts import build_rag_prompt
from app.services.retrieval.orchestrator import RetrievalOrchestrator

logger = logging.getLogger("ukweli.api.query")
router = APIRouter(tags=["RAG"])


@router.post("/query", response_model=QueryResponse)
async def query_rag(
    request: QueryRequest,
    db: AsyncSession = Depends(get_db_session),
    auth: AuthContext = Depends(require_auth),
    retriever: RetrievalOrchestrator = Depends(get_retrieval_orchestrator),
    llm: LLMGateway = Depends(get_llm_gateway),
    citation_fmt: CitationFormatter = Depends(get_citation_formatter),
) -> QueryResponse:
    """
    Full RAG pipeline:
    1. Validate auth + rate limit (multi-tier)
    2. Safety guardrails on query
    3. Hybrid retrieval (vector + keyword → RRF → reranker)
    4. Build LLM prompt with retrieved context
    5. Generate answer via HuggingFace Inference API
    6. Format citations with page/paragraph precision
    7. Log to audit trail
    8. Return response
    """
    start_time = time.monotonic()
    query_id = uuid.uuid4()

    # --- Step 1: Safety check ---
    safety_result = check_query_safety(request.query)
    if safety_result.blocked:
        latency_ms = int((time.monotonic() - start_time) * 1000)
        return QueryResponse(
            query_id=query_id,
            answer=safety_result.message,
            citations=[],
            retrieval_metadata=RetrievalMetadata(
                chunks_considered=0, latency_ms=latency_ms, graph_entities_used=[]
            ),
            confidence=Confidence.INSUFFICIENT_CONTEXT,
            suggested_followups=[],
        )

    # --- Step 2: Retrieve context ---
    retrieval_result = await retriever.retrieve(
        query=request.query,
        language=request.language,
        filters=request.filters,
        tier=auth.tier,
    )

    # --- Step 3: Build prompt and generate ---
    prompt_messages = build_rag_prompt(
        query=request.query,
        context_chunks=retrieval_result.context_blocks,
        language=request.language,
        mode=request.mode.value,
    )

    llm_response = await llm.generate(messages=prompt_messages)

    # --- Step 4: Format citations ---
    formatted = citation_fmt.format_response(
        raw_answer=llm_response.text,
        retrieved_chunks=retrieval_result.context_blocks,
    )

    latency_ms = int((time.monotonic() - start_time) * 1000)

    # --- Step 5: Determine confidence ---
    confidence = _assess_confidence(
        num_chunks=len(retrieval_result.context_blocks),
        top_score=retrieval_result.top_score,
    )

    # --- Step 6: Log to audit trail ---
    query_log = QueryLog(
        id=query_id,
        query_text=request.query,
        language=request.language,
        mode=request.mode.value,
        user_tier=auth.tier,
        user_id=auth.user.id if auth.user else None,
        api_key_id=auth.api_key.id if auth.api_key else None,
        fingerprint=auth.fingerprint,
        filters=request.filters.model_dump() if request.filters else None,
        answer=formatted.answer,
        citations=[c.model_dump() for c in formatted.citations] if formatted.citations else None,
        chunks_considered=retrieval_result.total_candidates,
        latency_ms=latency_ms,
        confidence=confidence.value,
        llm_model_used=llm_response.model_used,
        llm_tokens_used=llm_response.tokens_used,
        conversation_id=request.conversation_id,
    )
    db.add(query_log)
    await db.flush()

    logger.info(
        "Query processed: id=%s, tier=%s, latency=%dms, chunks=%d, confidence=%s",
        query_id,
        auth.tier,
        latency_ms,
        retrieval_result.total_candidates,
        confidence.value,
    )

    return QueryResponse(
        query_id=query_id,
        answer=formatted.answer,
        citations=formatted.citations,
        retrieval_metadata=RetrievalMetadata(
            chunks_considered=retrieval_result.total_candidates,
            latency_ms=latency_ms,
            graph_entities_used=retrieval_result.entities_used,
        ),
        confidence=confidence,
        suggested_followups=formatted.suggested_followups,
    )


def _assess_confidence(num_chunks: int, top_score: float) -> Confidence:
    """
    Heuristic confidence assessment based on retrieval quality.
    - HIGH: 3+ chunks with top score > 0.8
    - MEDIUM: 1+ chunks with top score > 0.5
    - LOW: chunks found but low scores
    - INSUFFICIENT: no chunks found
    """
    if num_chunks == 0:
        return Confidence.INSUFFICIENT_CONTEXT
    if num_chunks >= 3 and top_score > 0.8:
        return Confidence.HIGH
    if num_chunks >= 1 and top_score > 0.5:
        return Confidence.MEDIUM
    return Confidence.LOW