"""
title: Moderation Filter
author: justinh-rahb
author_url: https://github.com/justinh-rahb
funding_url: https://github.com/open-webui
version: 0.1.2
license: MIT
"""

from typing import Optional
from pydantic import BaseModel, Field
import requests
import os

# See for self-hosted moderations API setup:  https://github.com/matatonic/openedai-moderations

from utils.misc import get_last_user_message


class Filter:

    class Valves(BaseModel):
        priority: int = Field(
            default=0, description="Priority level for the filter operations."
        )
        moderation_api_base_url: str = Field(
            default="https://api.openai.com/v1",
            description="Base URL for the moderation API.",
        )
        openai_api_key: str = Field(
            default="", description="API key for OpenAI services."
        )
        moderation_model: str = Field(
            default="text-moderation-latest",
            description="Model to use for content moderation.",
        )
        enabled_for_admins: bool = Field(
            default=True, description="Whether moderation is enabled for admin users."
        )

    def __init__(self) -> None:
        self.valves = self.Valves()

    async def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        if __user__ is not None and (
            not __user__.get("role") == "admin" or self.valves.enabled_for_admins
        ):
            user_message = get_last_user_message(body["messages"])

            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.valves.openai_api_key}",
            }
            payload = {
                "model": self.valves.moderation_model,
                "input": user_message,
            }

            try:
                r = requests.post(
                    url=f"{self.valves.moderation_api_base_url}/moderations",
                    json=payload,
                    headers=headers,
                )

                r.raise_for_status()

                response = r.json()
            except Exception as e:
                print(f"Error: {e}")
                return body

            if response["results"][0]["flagged"]:
                flagged = [
                    k for k, v in response["results"][0]["categories"].items() if v
                ]
                print(f"Flagged: {flagged}")
                raise Exception(
                    f"Message contains flagged content: {', '.join(flagged)}"
                )

        return body

    def outlet(self, body: dict, __user__: Optional[dict] = None) -> dict:
        print(f"outlet:{__name__}")
        print(f"outlet:body:{body}")
        print(f"outlet:user:{__user__}")
        return body