"""
title: Content Violation Filter
author: jannikstdl
date: 2024-10-13
version: 1.0
description: A filter function that checks messages for content violations using Ollama moderation models 'llama-guard3' or 'shieldgemma' and flags them accordingly.
ShieldGemma: https://ollama.com/library/shieldgemma
LlamaGuard3: https://ollama.com/library/llama-guard3
"""

from typing import Optional, Callable, Any, Awaitable
from pydantic import BaseModel, Field
import requests
from open_webui.utils.misc import get_last_user_message


class Filter:
    class Valves(BaseModel):
        priority: int = Field(
            default=0, description="Priority level for the filter operations."
        )
        ollama_api_base_url: str = Field(
            default="http://localhost:11434",
            description="Base URL for the Ollama API.",
        )
        moderation_model: str = Field(
            default="llama-guard3:latest",
            description=(
                "Model to use for content moderation. Options are "
                "'llama-guard3' and 'shieldgemma'."
            ),
        )
        enabled_for_admins: bool = Field(
            default=True, description="Whether moderation is enabled for admin users."
        )
        enable_logs: bool = Field(
            default=False, description="Whether to print logs to console (Error/Info)."
        )

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

    def log(self, message: str, log_type: str = "INFO") -> None:
        if self.valves.enable_logs:
            if log_type == "ERROR":
                # ANSI code for light red text
                print(f"\033[91m{log_type}: {message}\033[0m")
            else:
                print(f"{log_type}: {message}")

    async def check_content(
        self,
        content: str,
        __event_emitter__: Callable[[Any], Awaitable[None]],
    ) -> None:
        headers = {
            "Content-Type": "application/json",
        }
        payload = {
            "model": self.valves.moderation_model,
            "messages": [{"role": "user", "content": content}],
            "stream": False,
        }
        try:
            r = requests.post(
                url=f"{self.valves.ollama_api_base_url}/api/chat",
                json=payload,
                headers=headers,
            )
            r.raise_for_status()
            response = r.json()
        except Exception as e:
            error_message = str(e)
            status_message = f"Content Moderation Filter: ERROR - Check your filter configuration or see logs if enabled in filter options for more infos."
            self.log(f"Content Moderation Filter: Error: {error_message}", "ERROR")
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {
                        "description": f"{status_message}",
                        "done": True,
                    },
                }
            )
            return

        moderation_result = response["message"]["content"].strip().lower()

        if self.valves.moderation_model.startswith("llama-guard3"):
            if moderation_result.startswith("unsafe"):
                categories = moderation_result.split("\n")[1:]
                classes = {
                    "s1": "Violent Crimes",
                    "s2": "Non-Violent Crimes",
                    "s3": "Sex-Related Crimes",
                    "s4": "Child Sexual Exploitation",
                    "s5": "Defamation",
                    "s6": "Specialized Advice",
                    "s7": "Privacy",
                    "s8": "Intellectual Property",
                    "s9": "Indiscriminate Weapons",
                    "s10": "Hate",
                    "s11": "Suicide & Self-Harm",
                    "s12": "Sexual Content",
                    "s13": "Elections",
                }
                category_descriptions = [
                    classes.get(cat.strip(), cat) for cat in categories
                ]
                self.log(
                    f"### Content Moderation Filter ###\n"
                    f"--- A message was flagged as content violation ---\n"
                    f"User Message: {content}\n"
                    f"Model: {self.valves.moderation_model}\n"
                    f"Flagged Categories: {category_descriptions}"
                )
                raise Exception(
                    f"Your request has been denied due to a content violation: {', '.join(category_descriptions)}"
                )
        elif self.valves.moderation_model.startswith("shieldgemma"):
            if moderation_result == "yes":
                self.log(
                    f"### Content Moderation Filter ###\n"
                    f"--- A message was flagged as content violation ---\n"
                    f"User Message: {content}\n"
                    f"Model: {self.valves.moderation_model}"
                )
                raise Exception(
                    "Your request has been denied due to a content violation"
                )

    async def inlet(
        self,
        body: dict,
        __user__: Optional[dict] = None,
        __event_emitter__: Callable[[Any], Awaitable[None]] = 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"])
            await self.check_content(user_message, __event_emitter__)
        return body