app/api/routes/classify.py

"""
POST /classify — document text classification endpoint.
"""

from fastapi import APIRouter, Depends, Header, HTTPException

from app.api.dependencies import verify_internal_service
from app.models.schemas import ClassifyRequest, ClassifyResponse, ErrorResponse
from app.ml.serving.registry import load_active_model
from app.ml.serving.inference import run_prediction
from app.core.logging import get_logger

logger = get_logger(__name__)

router = APIRouter(prefix="/analyze-injection", tags=["Prompt Injection Analysis"])


@router.post(
    "",
    response_model=ClassifyResponse,
    summary="Analyze document text for prompt injection threats",
    description=(
        "Accepts extracted text (from the Node.js PDF/OCR layer) "
        "and returns a risk label (safe/suspicious/injection) and confidence score."
    ),
    responses={
        401: {"model": ErrorResponse, "description": "Unauthorized — Missing or invalid X-Internal-Token header"},
        403: {"model": ErrorResponse, "description": "Forbidden — Client IP banned due to 3 failed token attempts"},
        422: {"model": ErrorResponse, "description": "Unprocessable Entity — Missing required fields or forbidden legacy keys"},
        503: {"model": ErrorResponse, "description": "Service Unavailable — Insufficient text (<5 words) or ML model load failure"},
    },
)
async def classify(
    req: ClassifyRequest,
):
    """Run the active RETVec+CNN model on fullText."""
    words = req.fullText.strip().split() if req.fullText else []
    if len(words) < 5:
        raise HTTPException(
            status_code=503,
            detail="insufficient_text"
        )

    try:
        model = await load_active_model()
    except Exception as e:
        logger.error("Classification model unavailable: %s", str(e))
        raise HTTPException(
            status_code=503,
            detail={"error": "Classification model unavailable", "detail": str(e)}
        )

    doc_id = req.documentId or "N/A"
    try:
        label, confidence = run_prediction(model, req.fullText)
    except Exception as e:
        logger.error("Inference prediction error for document %s: %s", doc_id, str(e), exc_info=True)
        raise HTTPException(
            status_code=500,
            detail=f"Inference failed: {str(e)}"
        )

    logger.info(
        "Classified document %s (length: %d chars, words: %d) → %s (confidence: %.2f)",
        doc_id,
        len(req.fullText),
        len(words),
        label,
        confidence,
    )

    return ClassifyResponse(
        label=label, confidence=confidence
    )