File size: 2,680 Bytes
215f97f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """
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
)
|