backend/app/api/query_router.py

"""
Query endpoint — the main pipeline entry point.

POST /api/v1/query         — run an agentic RAG query (authenticated)
GET  /api/v1/query/history — paginated query history for the current user
"""
import asyncio
import json
import logging
from datetime import datetime

from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from concurrent.futures import ThreadPoolExecutor

from app.agent.agent import run_agent
from app.auth import get_current_user
from app.config import get_settings
from app.database import get_db
from app.evaluation.trulens_eval import EvalScores, evaluate_async
from app.models import EvaluationResult, QueryLog, User
from app.visualization.viz_agent import VizCacheEntry, VIZ_CACHE, run_viz_agent

_viz_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="viz-agent")

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/query", tags=["query"])


# ── Pydantic schemas ──────────────────────────────────────────────────────────

class QueryRequest(BaseModel):
    query: str
    include_visualization: bool = False


class Citation(BaseModel):
    index: int
    title: str
    source_url: str
    full_citation: str


class QueryResponse(BaseModel):
    id: str
    query: str
    answer: str
    citations: list[Citation]
    chart_data: dict | None
    model_provider: str
    agent_steps: int
    created_at: datetime


class QueryHistoryItem(BaseModel):
    id: str
    query_text: str
    response_text: str | None
    model_provider: str | None
    agent_steps: int | None
    created_at: datetime


class EvalScoresPayload(BaseModel):
    relevance_score: float | None
    groundedness_score: float | None
    answer_relevance_score: float | None


class EvalStatusResponse(BaseModel):
    status: str  # "pending" | "complete"
    scores: EvalScoresPayload | None = None


# ── Helper ─────────────────────────────────────────────────────────────────────

def _persist_evaluation(
    eval_scores: EvalScores,
    query_log_id: str,
    loop: asyncio.AbstractEventLoop,
    db_session_factory,
) -> None:
    """
    Background callback: write TruLens scores to evaluation_results table.
    Runs when the evaluation future resolves.
    """
    async def _write():
        async with db_session_factory() as session:
            row = EvaluationResult(
                query_log_id=query_log_id,
                relevance_score=eval_scores.relevance_score,
                groundedness_score=eval_scores.groundedness_score,
                answer_relevance_score=eval_scores.answer_relevance_score,
                trulens_record_id=eval_scores.trulens_record_id,
            )
            session.add(row)
            await session.commit()

    asyncio.run_coroutine_threadsafe(_write(), loop)


# ── Endpoints ─────────────────────────────────────────────────────────────────

@router.post("", response_model=QueryResponse)
async def submit_query(
    body: QueryRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    settings = get_settings()

    if not body.query.strip():
        raise HTTPException(status_code=422, detail="Query must not be empty")

    # ── Run the agent (synchronous work dispatched to thread pool) ────────────
    loop = asyncio.get_event_loop()
    agent_result: dict = await loop.run_in_executor(
        None, run_agent, current_user.id, body.query, settings
    )

    answer: str = agent_result.get("answer", "")
    citations_raw: list = agent_result.get("citations", [])
    chart_data: dict | None = agent_result.get("chart_data")
    agent_steps: int = agent_result.get("agent_steps", 0)

    # ── Persist query log ─────────────────────────────────────────────────────
    log_entry = QueryLog(
        user_id=current_user.id,
        query_text=body.query,
        response_text=answer,
        model_provider=settings.model_provider,
        agent_steps=agent_steps,
    )
    db.add(log_entry)
    await db.commit()
    await db.refresh(log_entry)

    # ── Schedule TruLens evaluation asynchronously ────────────────────────────
    # Prefer authoritative pgvector_retriever output (reranked top-5 + chunk text);
    # JSON "context" may be truncated; citation-only fallback lacks Content:.
    citation_fallback = "\n\n".join(
        c.get("full_citation", "") + "\n" + c.get("title", "")
        for c in citations_raw
    )
    context_text = (
        agent_result.get("reranked_context", "")
        or agent_result.get("context", "")
        or citation_fallback
    )
    if not agent_result.get("reranked_context") and citations_raw:
        logger.debug(
            "TruLens context: empty reranked_context but citations present — using JSON context or citation fallback"
        )

    eval_future = await evaluate_async(
        query=body.query,
        context=context_text,
        answer=answer,
        query_log_id=log_entry.id,
        settings=settings,
    )

    # Wire up persistence callback when evaluation completes (fire-and-forget)
    from app.database import AsyncSessionLocal

    def _on_eval_done(future):
        try:
            scores: EvalScores = future.result()
            _persist_evaluation(scores, log_entry.id, loop, AsyncSessionLocal)
        except Exception as exc:
            logger.error("Evaluation persistence failed: %s", exc)

    eval_future.add_done_callback(_on_eval_done)

    # Normalise citations into typed objects
    citations = [
        Citation(
            index=c.get("index", i + 1),
            title=c.get("title", ""),
            source_url=c.get("source_url", ""),
            full_citation=c.get("full_citation", ""),
        )
        for i, c in enumerate(citations_raw)
    ]

    # ── Fire visualization agent (fire-and-forget, same pattern as TruLens) ────
    if body.include_visualization:
        # Seed a "pending" entry immediately so the poll endpoint returns 200
        VIZ_CACHE[log_entry.id] = VizCacheEntry(status="pending")
        _viz_executor.submit(run_viz_agent, log_entry.id, body.query, answer, settings)
        logger.debug("Visualization agent submitted for query_id=%s", log_entry.id)

    return QueryResponse(
        id=log_entry.id,
        query=body.query,
        answer=answer,
        citations=citations,
        chart_data=chart_data,
        model_provider=settings.model_provider,
        agent_steps=agent_steps,
        created_at=log_entry.created_at,
    )


@router.get("/{query_id}/evaluation", response_model=EvalStatusResponse)
async def get_evaluation(
    query_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """
    Poll for TruLens RAG Triad scores for a completed query.
    Returns {"status": "pending"} while the background eval thread is still running,
    and {"status": "complete", "scores": {...}} once the row is persisted.
    """
    # Verify query exists and belongs to this user
    log_result = await db.execute(
        select(QueryLog).where(
            QueryLog.id == query_id,
            QueryLog.user_id == current_user.id,
        )
    )
    if log_result.scalar_one_or_none() is None:
        raise HTTPException(status_code=404, detail="Query not found")

    eval_result = await db.execute(
        select(EvaluationResult).where(EvaluationResult.query_log_id == query_id)
    )
    row = eval_result.scalar_one_or_none()

    if row is None:
        return EvalStatusResponse(status="pending")

    return EvalStatusResponse(
        status="complete",
        scores=EvalScoresPayload(
            relevance_score=float(row.relevance_score) if row.relevance_score is not None else None,
            groundedness_score=float(row.groundedness_score) if row.groundedness_score is not None else None,
            answer_relevance_score=float(row.answer_relevance_score) if row.answer_relevance_score is not None else None,
        ),
    )


@router.get("/history", response_model=list[QueryHistoryItem])
async def query_history(
    limit: int = 20,
    offset: int = 0,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    result = await db.execute(
        select(QueryLog)
        .where(QueryLog.user_id == current_user.id)
        .order_by(QueryLog.created_at.desc())
        .limit(limit)
        .offset(offset)
    )
    logs = result.scalars().all()
    return [
        QueryHistoryItem(
            id=log.id,
            query_text=log.query_text,
            response_text=log.response_text,
            model_provider=log.model_provider,
            agent_steps=log.agent_steps,
            created_at=log.created_at,
        )
        for log in logs
    ]