src/agents/ModerationAgent.py

import logging
from typing import Any, Dict

from dotenv import load_dotenv
from openai import OpenAI

logger = logging.getLogger(__name__)


class ModerationAgent:
    """
    Wrapper around OpenAI Moderation API.

    Uses the latest moderation model and returns a simple flagged/not-flagged result.
    """

    def __init__(self):
        load_dotenv()
        self.client = OpenAI()

    def classify(self, text: str) -> Dict[str, Any]:
        text = (text or "").strip()
        if not text:
            return {"flagged": False, "categories": {}, "source": "empty"}

        try:
            moderation = self.client.moderations.create(
                model="omni-moderation-latest",
                input=text,
            )
            result = moderation.results[0]
            categories = (
                result.categories.model_dump()
                if hasattr(result.categories, "model_dump")
                else dict(result.categories)
            )
            category_scores = (
                result.category_scores.model_dump()
                if hasattr(result.category_scores, "model_dump")
                else dict(result.category_scores)
            )
            flagged = bool(result.flagged)
            logger.info(
                "ModerationAgent: flagged=%s categories=%s",
                flagged,
                [k for k, v in categories.items() if v],
            )
            return {
                "flagged": flagged,
                "categories": categories,
                "category_scores": category_scores,
                "source": "openai_moderation",
            }
        except Exception as e:
            logger.exception("Moderation API call failed")
            return {
                "flagged": False,
                "categories": {},
                "error": str(e),
                "source": "error",
            }

    def is_flagged(self, text: str) -> bool:
        return bool(self.classify(text).get("flagged"))