openai_moderation.py — langchain Source File

Source Code

"""Agent middleware that integrates OpenAI's moderation endpoint."""

from __future__ import annotations

import json
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Literal, cast

from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
from openai import AsyncOpenAI, OpenAI
from openai.types import Moderation, ModerationModel

if TYPE_CHECKING:  # pragma: no cover
    from langgraph.runtime import Runtime

ViolationStage = Literal["input", "output", "tool"]

DEFAULT_VIOLATION_TEMPLATE = (
    "I'm sorry, but I can't comply with that request. It was flagged for {categories}."
)


class OpenAIModerationError(RuntimeError):
    """Raised when OpenAI flags content and `exit_behavior` is set to ``"error"``."""

    def __init__(
        self,
        *,
        content: str,
        stage: ViolationStage,
        result: Moderation,
        message: str,
    ) -> None:
        """Initialize the error with violation details.

        Args:
            content: The content that was flagged.
            stage: The stage where the violation occurred.
            result: The moderation result from OpenAI.
            message: The error message.
        """
        super().__init__(message)
        self.content = content
        self.stage = stage
        self.result = result


class OpenAIModerationMiddleware(AgentMiddleware[AgentState[Any], Any]):
    """Moderate agent traffic using OpenAI's moderation endpoint."""

    def __init__(
        self,
        *,
        model: ModerationModel = "omni-moderation-latest",
        check_input: bool = True,
        check_output: bool = True,
        check_tool_results: bool = False,
        exit_behavior: Literal["error", "end", "replace"] = "end",
        violation_message: str | None = None,
// ... (425 more lines)