src/api/qa.py

"""
QA endpoint for grant analyst queries.

Clean implementation using the QA service layer with validated schemas.
"""
import asyncio
import json
from typing import Dict, Any
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from uuid import uuid4

from src.logging.logger import get_logger
from analyzer.models import QARequest, QAChunk
from analyzer.qa_service import stream_answer, answer_question as qa_answer

logger = get_logger()
router = APIRouter(prefix="/qa", tags=["qa"])


@router.post("/")
async def qa_endpoint(request: QARequest) -> Dict[str, Any]:
    """
    Answer a grant-related question (non-streaming).

    Uses the QA service layer with validated models and prompt injection hardening.

    Args:
        request: QARequest with query and optional filters

    Returns:
        Dict with answer, citations, and metadata

    Raises:
        HTTPException: On service failures
    """
    # Ensure session ID is set
    if not request.session_id:
        request.session_id = str(uuid4())

    try:
        # Use service layer
        result = qa_answer(request)

        # Log interaction
        logger.info(
            f"QA request processed: session={request.session_id}, "
            f"latency={result['latency_ms']}ms, success={result['success']}"
        )

        return result

    except Exception as e:
        logger.error(f"QA endpoint error: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/stream")
async def qa_stream_endpoint(request: QARequest) -> StreamingResponse:
    """
    Answer a grant-related question with Server-Sent Events streaming.

    Streams the LLM response as it generates for better UX.

    Args:
        request: QARequest with query and optional filters

    Returns:
        StreamingResponse with SSE-formatted NDJSON stream

    Raises:
        HTTPException: On service failures
    """
    # Ensure session ID is set
    if not request.session_id:
        request.session_id = str(uuid4())

    async def event_stream():
        """Generate SSE event stream."""
        try:
            # Stream from service layer
            for chunk in stream_answer(request):
                # Serialize chunk to NDJSON
                chunk_json = chunk.model_dump_json()
                yield f"data: {chunk_json}\n\n"

                # Allow other tasks to run
                await asyncio.sleep(0)

        except Exception as e:
            logger.error(f"Stream error: {e}", exc_info=True)
            # Send error chunk
            error_chunk = QAChunk(type="error", error=str(e))
            yield f"data: {error_chunk.model_dump_json()}\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"  # Disable nginx buffering
        }
    )


@router.get("/test")
async def test_qa() -> Dict[str, str]:
    """Test endpoint to verify QA service is running."""
    return {
        "status": "ok",
        "message": "QA service is running",
        "endpoint": "/qa"
    }