api.py

import asyncio
import json
import logging
import os
import time
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from fastapi import BackgroundTasks, Depends, FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.security import APIKeyHeader

from caching import CacheTier, cache_manager
from graph import app as crag_app
from graph import build_initial_state
from models import QueryRequest, QueryResponse

# Configure structured runtime logging
logger = logging.getLogger("k_state.api")
logging.basicConfig(level=logging.INFO)

API_KEY_HEADER = APIKeyHeader(name="X-API-KEY", auto_error=False)

# 1. Defend Against Insecure Deployments: Enforce environment variables explicitly on initialization
SECRET_TOKEN = os.getenv("CRAG_SECRET_TOKEN")
if not SECRET_TOKEN:
    raise RuntimeError(
        "CRITICAL ERROR: 'CRAG_SECRET_TOKEN' environment variable is missing. "
        "Server startup terminated to prevent unauthorized fallback vulnerabilities."
    )


async def verify_api_key(api_key: str = Depends(API_KEY_HEADER)):
    """FastAPI Dependency Injection credential verification gate.
    
    4. Enforce Strict Authentication: Rejects missing (None) or mismatched keys uniformly.
    """
    if api_key != SECRET_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, 
            detail="Invalid or missing API security credentials."
        )


@asynccontextmanager
async def lifespan(app: FastAPI):
    yield


app = FastAPI(title="CRAG Async API Server", version="3.0.0", lifespan=lifespan)

# 5. Correct CORS Policies: Disabled credential passage to align cleanly with wildcard restrictions
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


# --- Secure Global Exception Mapping Handlers ---

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    return JSONResponse(
        status_code=exc.status_code, 
        content={"error": "ClientRequestError", "detail": exc.detail}
    )


@app.exception_handler(Exception)
async def global_generic_exception_handler(request: Request, exc: Exception):
    """3. Eliminate Exception Leakage: Mask raw stack traces from client exposure while logging internally."""
    logger.error(f"UNHANDLED BACKGROUND SYSTEM EXCEPTION: {str(exc)}", exc_info=True)
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"error": "InternalServerError"}
    )


@app.get("/")
async def root():
    return {"service": "CRAG Async API Engine", "status": "healthy", "cache_enabled": cache_manager.enabled}


@app.post("/query", response_model=QueryResponse, dependencies=[Depends(verify_api_key)])
async def query(request: QueryRequest, background_tasks: BackgroundTasks):
    # 2. Precision Latency Telemetry: Utilizing counter clocks for monotonic precision
    start_time = time.perf_counter()
    
    if not request.question or len(request.question.strip()) < 3:
        raise HTTPException(status_code=400, detail="Question must be at least 3 characters")

    if request.use_cache and cache_manager.enabled:
        cached = cache_manager.get(request.question)
        if cached:
            return QueryResponse(
                question=request.question,
                answer=cached.answer,
                verdict=cached.verdict,
                reason="Cached response",
                agentic_turns=cached.agentic_turns,
                cached=True,
                latency_ms=(time.perf_counter() - start_time) * 1000,
            )

    # Execute workflow graph asynchronously on the single-thread event loop
    result = await crag_app.ainvoke(build_initial_state(request.question))

    response = QueryResponse(
        question=request.question,
        answer=result.get("answer", "No answer generated"),
        verdict=result.get("verdict", "INCORRECT"),
        reason=result.get("reason", ""),
        agentic_turns=result.get("agentic_turns", 0),
        cached=False,
        latency_ms=(time.perf_counter() - start_time) * 1000,
        tokens_used=result.get("total_tokens", 0),
    )

    # Save to cache tier via a non-blocking background task worker
    if request.use_cache and cache_manager.enabled:
        background_tasks.add_task(cache_manager.set, request.question, response, CacheTier.HOT)

    return response


@app.post("/query/stream", dependencies=[Depends(verify_api_key)])
async def query_stream(request: QueryRequest):
    """Asynchronous Streaming endpoint using native Server-Sent Events (SSE)."""
    
    async def event_generator() -> AsyncGenerator[str, None]:
        def emit(event_type: str, data: dict) -> str:
            return f"data: {json.dumps({'type': event_type, **data})}\n\n"

        yield emit("start", {})
        initial_state = build_initial_state(request.question)
        
        try:
            # Native async streaming iteration over active LangGraph graph updates
            async for chunk in crag_app.astream(initial_state, stream_mode="updates"):
                for node_name, node_update in chunk.items():
                    yield emit("stage", {"node": node_name})
                    if "verdict" in node_update:
                        yield emit("verdict", {"verdict": node_update["verdict"]})
                    if "answer" in node_update:
                        yield emit("result", {"answer": node_update["answer"]})
            yield emit("end", {})
        except Exception as err:
            logger.error(f"STREAM PROCESSING ERROR: {str(err)}", exc_info=True)
            yield emit("error", {})  # 3. Mask trace data in streaming channels

    return StreamingResponse(event_generator(), media_type="text/event-stream")