detector.py

"""
detector.py

LLM-based risk classifier.
"""

import json
from dataclasses import dataclass
from logger import logger
from llm_client import call_llm
from prompts import build_detector_messages


VALID_CATEGORIES = {
    "SAFE",
    "LOW_RISK",
    "MEDIUM_RISK",
    "HIGH_RISK"
}

@dataclass
class DetectionResult:
    category: str
    confidence: float
    details: str

def _strip_markdown(text: str) -> str:
    """
    Remove markdown code fences if the model wrapped JSON in them.
    """
    text = text.strip()

    if not text.startswith("```"):
        return text

    lines = text.splitlines()

    # Drop opening fence (may be ``` or ```json)
    if lines[0].strip().startswith("```"):
        lines = lines[1:]

    # Drop closing fence
    if lines and lines[-1].strip() == "```":
        lines = lines[:-1]

    return "\n".join(lines).strip()

def detect(message: str) -> DetectionResult:
    """
    Classify a user message for risk level using the LLM.
    Falls back to SAFE on any error.
    """
    messages = build_detector_messages(message)
    raw_response = call_llm(messages)

    try:
        clean_text = _strip_markdown(raw_response)
        data = json.loads(clean_text)

        category = data.get("category", "SAFE")
        if category not in VALID_CATEGORIES:
            category = "SAFE"

        confidence = float(data.get("confidence", 0.0))
        confidence = max(0.0, min(1.0, confidence))

        details = data.get("reason", "No explanation provided.")

        result = DetectionResult(
            category=category,
            confidence=confidence,
            details=details
        )

        logger.info(
            "Detector | category=%s confidence=%.2f",
            result.category,
            result.confidence
        )

        return result

    except (json.JSONDecodeError, ValueError, TypeError, KeyError) as error:
        logger.exception("Detector parse failed for message: %s", message)
        return DetectionResult(
            category="SAFE",
            confidence=0.0,
            details=f"Detection failed — defaulted to SAFE. ({error})"
        )