coco-framework 0.1.0__py3-none-any.whl

coco_openai/middleware/__init__.py

@@ -0,0 +1,11 @@
+"""Middleware implementations for OpenAI-backed agents."""
+
+from coco_openai.middleware.openai_moderation import (
+    OpenAIModerationError,
+    OpenAIModerationMiddleware,
+)
+
+__all__ = [
+    "OpenAIModerationError",
+    "OpenAIModerationMiddleware",
+]

coco_openai/middleware/openai_moderation.py

@@ -0,0 +1,484 @@
+"""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 coco.agents.middleware.types import AgentMiddleware, AgentState, hook_config
+from coco_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 coco_graph.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,
+        client: OpenAI | None = None,
+        async_client: AsyncOpenAI | None = None,
+    ) -> None:
+        """Create the middleware instance.
+
+        Args:
+            model: OpenAI moderation model to use.
+            check_input: Whether to check user input messages.
+            check_output: Whether to check model output messages.
+            check_tool_results: Whether to check tool result messages.
+            exit_behavior: How to handle violations
+                (`'error'`, `'end'`, or `'replace'`).
+            violation_message: Custom template for violation messages.
+            client: Optional pre-configured OpenAI client to reuse.
+                If not provided, a new client will be created.
+            async_client: Optional pre-configured AsyncOpenAI client to reuse.
+                If not provided, a new async client will be created.
+        """
+        super().__init__()
+        self.model = model
+        self.check_input = check_input
+        self.check_output = check_output
+        self.check_tool_results = check_tool_results
+        self.exit_behavior = exit_behavior
+        self.violation_message = violation_message
+
+        self._client = client
+        self._async_client = async_client
+
+    @hook_config(can_jump_to=["end"])
+    def before_model(
+        self, state: AgentState[Any], runtime: Runtime[Any]
+    ) -> dict[str, Any] | None:  # type: ignore[override]
+        """Moderate user input and tool results before the model is called.
+
+        Args:
+            state: Current agent state containing messages.
+            runtime: Agent runtime context.
+
+        Returns:
+            Updated state with moderated messages, or `None` if no changes.
+        """
+        if not self.check_input and not self.check_tool_results:
+            return None
+
+        messages = list(state.get("messages", []))
+        if not messages:
+            return None
+
+        return self._moderate_inputs(messages)
+
+    @hook_config(can_jump_to=["end"])
+    def after_model(
+        self, state: AgentState[Any], runtime: Runtime[Any]
+    ) -> dict[str, Any] | None:  # type: ignore[override]
+        """Moderate model output after the model is called.
+
+        Args:
+            state: Current agent state containing messages.
+            runtime: Agent runtime context.
+
+        Returns:
+            Updated state with moderated messages, or `None` if no changes.
+        """
+        if not self.check_output:
+            return None
+
+        messages = list(state.get("messages", []))
+        if not messages:
+            return None
+
+        return self._moderate_output(messages)
+
+    @hook_config(can_jump_to=["end"])
+    async def abefore_model(
+        self, state: AgentState[Any], runtime: Runtime[Any]
+    ) -> dict[str, Any] | None:  # type: ignore[override]
+        """Async version of before_model.
+
+        Args:
+            state: Current agent state containing messages.
+            runtime: Agent runtime context.
+
+        Returns:
+            Updated state with moderated messages, or `None` if no changes.
+        """
+        if not self.check_input and not self.check_tool_results:
+            return None
+
+        messages = list(state.get("messages", []))
+        if not messages:
+            return None
+
+        return await self._amoderate_inputs(messages)
+
+    @hook_config(can_jump_to=["end"])
+    async def aafter_model(
+        self, state: AgentState[Any], runtime: Runtime[Any]
+    ) -> dict[str, Any] | None:  # type: ignore[override]
+        """Async version of after_model.
+
+        Args:
+            state: Current agent state containing messages.
+            runtime: Agent runtime context.
+
+        Returns:
+            Updated state with moderated messages, or `None` if no changes.
+        """
+        if not self.check_output:
+            return None
+
+        messages = list(state.get("messages", []))
+        if not messages:
+            return None
+
+        return await self._amoderate_output(messages)
+
+    def _moderate_inputs(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        working = list(messages)
+        modified = False
+
+        if self.check_tool_results:
+            action = self._moderate_tool_messages(working)
+            if action:
+                if "jump_to" in action:
+                    return action
+                working = cast("list[BaseMessage]", action["messages"])
+                modified = True
+
+        if self.check_input:
+            action = self._moderate_user_message(working)
+            if action:
+                if "jump_to" in action:
+                    return action
+                working = cast("list[BaseMessage]", action["messages"])
+                modified = True
+
+        if modified:
+            return {"messages": working}
+
+        return None
+
+    async def _amoderate_inputs(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        working = list(messages)
+        modified = False
+
+        if self.check_tool_results:
+            action = await self._amoderate_tool_messages(working)
+            if action:
+                if "jump_to" in action:
+                    return action
+                working = cast("list[BaseMessage]", action["messages"])
+                modified = True
+
+        if self.check_input:
+            action = await self._amoderate_user_message(working)
+            if action:
+                if "jump_to" in action:
+                    return action
+                working = cast("list[BaseMessage]", action["messages"])
+                modified = True
+
+        if modified:
+            return {"messages": working}
+
+        return None
+
+    def _moderate_output(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        last_ai_idx = self._find_last_index(messages, AIMessage)
+        if last_ai_idx is None:
+            return None
+
+        ai_message = messages[last_ai_idx]
+        text = self._extract_text(ai_message)
+        if not text:
+            return None
+
+        result = self._moderate(text)
+        if not result.flagged:
+            return None
+
+        return self._apply_violation(
+            messages, index=last_ai_idx, stage="output", content=text, result=result
+        )
+
+    async def _amoderate_output(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        last_ai_idx = self._find_last_index(messages, AIMessage)
+        if last_ai_idx is None:
+            return None
+
+        ai_message = messages[last_ai_idx]
+        text = self._extract_text(ai_message)
+        if not text:
+            return None
+
+        result = await self._amoderate(text)
+        if not result.flagged:
+            return None
+
+        return self._apply_violation(
+            messages, index=last_ai_idx, stage="output", content=text, result=result
+        )
+
+    def _moderate_tool_messages(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        last_ai_idx = self._find_last_index(messages, AIMessage)
+        if last_ai_idx is None:
+            return None
+
+        working = list(messages)
+        modified = False
+
+        for idx in range(last_ai_idx + 1, len(working)):
+            msg = working[idx]
+            if not isinstance(msg, ToolMessage):
+                continue
+
+            text = self._extract_text(msg)
+            if not text:
+                continue
+
+            result = self._moderate(text)
+            if not result.flagged:
+                continue
+
+            action = self._apply_violation(
+                working, index=idx, stage="tool", content=text, result=result
+            )
+            if action:
+                if "jump_to" in action:
+                    return action
+                working = cast("list[BaseMessage]", action["messages"])
+                modified = True
+
+        if modified:
+            return {"messages": working}
+
+        return None
+
+    async def _amoderate_tool_messages(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        last_ai_idx = self._find_last_index(messages, AIMessage)
+        if last_ai_idx is None:
+            return None
+
+        working = list(messages)
+        modified = False
+
+        for idx in range(last_ai_idx + 1, len(working)):
+            msg = working[idx]
+            if not isinstance(msg, ToolMessage):
+                continue
+
+            text = self._extract_text(msg)
+            if not text:
+                continue
+
+            result = await self._amoderate(text)
+            if not result.flagged:
+                continue
+
+            action = self._apply_violation(
+                working, index=idx, stage="tool", content=text, result=result
+            )
+            if action:
+                if "jump_to" in action:
+                    return action
+                working = cast("list[BaseMessage]", action["messages"])
+                modified = True
+
+        if modified:
+            return {"messages": working}
+
+        return None
+
+    def _moderate_user_message(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        idx = self._find_last_index(messages, HumanMessage)
+        if idx is None:
+            return None
+
+        message = messages[idx]
+        text = self._extract_text(message)
+        if not text:
+            return None
+
+        result = self._moderate(text)
+        if not result.flagged:
+            return None
+
+        return self._apply_violation(
+            messages, index=idx, stage="input", content=text, result=result
+        )
+
+    async def _amoderate_user_message(
+        self, messages: Sequence[BaseMessage]
+    ) -> dict[str, Any] | None:
+        idx = self._find_last_index(messages, HumanMessage)
+        if idx is None:
+            return None
+
+        message = messages[idx]
+        text = self._extract_text(message)
+        if not text:
+            return None
+
+        result = await self._amoderate(text)
+        if not result.flagged:
+            return None
+
+        return self._apply_violation(
+            messages, index=idx, stage="input", content=text, result=result
+        )
+
+    def _apply_violation(
+        self,
+        messages: Sequence[BaseMessage],
+        *,
+        index: int | None,
+        stage: ViolationStage,
+        content: str,
+        result: Moderation,
+    ) -> dict[str, Any] | None:
+        violation_text = self._format_violation_message(content, result)
+
+        if self.exit_behavior == "error":
+            raise OpenAIModerationError(
+                content=content,
+                stage=stage,
+                result=result,
+                message=violation_text,
+            )
+
+        if self.exit_behavior == "end":
+            return {"jump_to": "end", "messages": [AIMessage(content=violation_text)]}
+
+        if index is None:
+            return None
+
+        new_messages = list(messages)
+        original = new_messages[index]
+        new_messages[index] = cast(
+            BaseMessage, original.model_copy(update={"content": violation_text})
+        )
+        return {"messages": new_messages}
+
+    def _moderate(self, text: str) -> Moderation:
+        if self._client is None:
+            self._client = self._build_client()
+        response = self._client.moderations.create(model=self.model, input=text)
+        return response.results[0]
+
+    async def _amoderate(self, text: str) -> Moderation:
+        if self._async_client is None:
+            self._async_client = self._build_async_client()
+        response = await self._async_client.moderations.create(
+            model=self.model, input=text
+        )
+        return response.results[0]
+
+    def _build_client(self) -> OpenAI:
+        self._client = OpenAI()
+        return self._client
+
+    def _build_async_client(self) -> AsyncOpenAI:
+        self._async_client = AsyncOpenAI()
+        return self._async_client
+
+    def _format_violation_message(self, content: str, result: Moderation) -> str:
+        # Convert categories to dict and filter for flagged items
+        categories_dict = result.categories.model_dump()
+        categories = [
+            name.replace("_", " ")
+            for name, flagged in categories_dict.items()
+            if flagged
+        ]
+        category_label = (
+            ", ".join(categories) if categories else "OpenAI's safety policies"
+        )
+        template = self.violation_message or DEFAULT_VIOLATION_TEMPLATE
+        scores_json = json.dumps(result.category_scores.model_dump(), sort_keys=True)
+        try:
+            message = template.format(
+                categories=category_label,
+                category_scores=scores_json,
+                original_content=content,
+            )
+        except KeyError:
+            message = template
+        return message
+
+    def _find_last_index(
+        self, messages: Sequence[BaseMessage], message_type: type[BaseMessage]
+    ) -> int | None:
+        for idx in range(len(messages) - 1, -1, -1):
+            if isinstance(messages[idx], message_type):
+                return idx
+        return None
+
+    def _extract_text(self, message: BaseMessage) -> str | None:
+        if message.content is None:
+            return None
+        text_accessor = getattr(message, "text", None)
+        if text_accessor is None:
+            return str(message.content)
+        text = str(text_accessor)
+        return text if text else None
+
+
+__all__ = [
+    "OpenAIModerationError",
+    "OpenAIModerationMiddleware",
+]

coco_openai/output_parsers/__init__.py

@@ -0,0 +1,9 @@
+"""Output parsers for OpenAI tools."""
+
+from coco_core.output_parsers.openai_tools import (
+    JsonOutputKeyToolsParser,
+    JsonOutputToolsParser,
+    PydanticToolsParser,
+)
+
+__all__ = ["JsonOutputKeyToolsParser", "JsonOutputToolsParser", "PydanticToolsParser"]

coco_openai/output_parsers/tools.py

@@ -0,0 +1,9 @@
+"""Output parsers for OpenAI tools."""
+
+from coco_core.output_parsers.openai_tools import (
+    JsonOutputKeyToolsParser,
+    JsonOutputToolsParser,
+    PydanticToolsParser,
+)
+
+__all__ = ["JsonOutputKeyToolsParser", "JsonOutputToolsParser", "PydanticToolsParser"]

coco_openai/tools/__init__.py

@@ -0,0 +1,5 @@
+"""Tools package for OpenAI integrations."""
+
+from coco_openai.tools.custom_tool import custom_tool
+
+__all__ = ["custom_tool"]

coco_openai/tools/custom_tool.py

@@ -0,0 +1,111 @@
+"""Custom tool decorator for OpenAI custom tools."""
+
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from coco_core.tools import tool
+
+
+def _make_wrapped_func(func: Callable[..., str]) -> Callable[..., list[dict[str, Any]]]:
+    def wrapped(x: str) -> list[dict[str, Any]]:
+        return [{"type": "custom_tool_call_output", "output": func(x)}]
+
+    return wrapped
+
+
+def _make_wrapped_coroutine(
+    coroutine: Callable[..., Awaitable[str]],
+) -> Callable[..., Awaitable[list[dict[str, Any]]]]:
+    async def wrapped(*args: Any, **kwargs: Any) -> list[dict[str, Any]]:
+        result = await coroutine(*args, **kwargs)
+        return [{"type": "custom_tool_call_output", "output": result}]
+
+    return wrapped
+
+
+def custom_tool(*args: Any, **kwargs: Any) -> Any:
+    """Decorator to create an OpenAI custom tool.
+
+    Custom tools allow for tools with (potentially long) freeform string inputs.
+
+    See below for an example using CocoGraph:
+
+    ```python
+    @custom_tool
+    def execute_code(code: str) -> str:
+        \"\"\"Execute python code.\"\"\"
+        return "27"
+
+
+    model = ChatOpenAI(model="gpt-5", output_version="responses/v1")
+
+    agent = create_react_agent(model, [execute_code])
+
+    input_message = {"role": "user", "content": "Use the tool to calculate 3^3."}
+    for step in agent.stream(
+        {"messages": [input_message]},
+        stream_mode="values",
+    ):
+        step["messages"][-1].pretty_print()
+    ```
+
+    You can also specify a format for a corresponding context-free grammar using the
+    `format` kwarg:
+
+    ```python
+    from coco_openai import ChatOpenAI, custom_tool
+    from coco_graph.prebuilt import create_react_agent
+
+    grammar = \"\"\"
+    start: expr
+    expr: term (SP ADD SP term)* -> add
+    | term
+    term: factor (SP MUL SP factor)* -> mul
+    | factor
+    factor: INT
+    SP: " "
+    ADD: "+"
+    MUL: "*"
+    %import common.INT
+    \"\"\"
+
+    format = {"type": "grammar", "syntax": "lark", "definition": grammar}
+
+    # highlight-next-line
+    @custom_tool(format=format)
+    def do_math(input_string: str) -> str:
+        \"\"\"Do a mathematical operation.\"\"\"
+        return "27"
+
+
+    model = ChatOpenAI(model="gpt-5", output_version="responses/v1")
+
+    agent = create_react_agent(model, [do_math])
+
+    input_message = {"role": "user", "content": "Use the tool to calculate 3^3."}
+    for step in agent.stream(
+        {"messages": [input_message]},
+        stream_mode="values",
+    ):
+        step["messages"][-1].pretty_print()
+    ```
+    """
+
+    def decorator(func: Callable[..., Any]) -> Any:
+        metadata = {"type": "custom_tool"}
+        if "format" in kwargs:
+            metadata["format"] = kwargs.pop("format")
+        tool_obj = tool(infer_schema=False, **kwargs)(func)
+        tool_obj.metadata = metadata
+        tool_obj.description = func.__doc__
+        if inspect.iscoroutinefunction(func):
+            tool_obj.coroutine = _make_wrapped_coroutine(func)
+        else:
+            tool_obj.func = _make_wrapped_func(func)
+        return tool_obj
+
+    if args and callable(args[0]) and not kwargs:
+        return decorator(args[0])
+
+    return decorator