coco-framework 0.1.0__py3-none-any.whl

coco/agents/middleware/internal_call_transformer.py

@@ -0,0 +1,108 @@
+"""Tag and filter middleware-internal model calls.
+
+Tag internal calls with `internal_call_metadata()` and declare
+`transformers = (InternalCallTransformer,)` on the middleware class to keep
+them out of `run.messages`. Both APIs are public for third-party middleware.
+"""
+
+from __future__ import annotations
+
+import secrets
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from coco_core.messages import BaseMessage
+from coco_graph.stream import StreamTransformer
+
+if TYPE_CHECKING:
+    from coco_graph.stream._types import ProtocolEvent
+
+INTERNAL_CALL_METADATA_KEY = "lc_internal_call"
+"""Metadata key marking a model call as internal to middleware.
+
+Kept separate from `lc_source` so filtering doesn't affect its existing
+semantics.
+"""
+
+_INTERNAL_CALL_TOKEN = secrets.token_hex(16)
+"""Process-local marker used to prevent callers from spoofing internal calls.
+
+A random token prevents user-supplied metadata from hiding real model calls
+from `run.messages`.
+"""
+
+
+def internal_call_metadata() -> dict[str, Any]:
+    """Return metadata that marks a model call as internal to middleware.
+
+    Returns:
+        A mapping to merge into a model call's `config["metadata"]`.
+    """
+    return {INTERNAL_CALL_METADATA_KEY: _INTERNAL_CALL_TOKEN}
+
+
+class InternalCallTransformer(StreamTransformer):
+    """Keep internal model calls out of `run.messages` and the raw event log.
+
+    Used by middleware that makes internal model calls and runs before built-in
+    transformers.
+
+    For tagged events, streamed `message-start` events are marked as tool-role and
+    whole-`AIMessage` payloads are cleared so `MessagesTransformer` ignores them.
+    The mutated events are then dropped from the raw log.
+
+    Only events within this transformer's scope are modified.
+    """
+
+    before_builtins: ClassVar[bool] = True
+    required_stream_modes: ClassVar[tuple[str, ...]] = ("messages",)
+
+    def __init__(self, scope: tuple[str, ...] = ()) -> None:
+        """Initialize the transformer, scoped to a single mux/namespace.
+
+        Args:
+            scope: The namespace tuple the owning mux is scoped to.
+        """
+        super().__init__(scope)
+        self._scope_list: list[str] = list(scope)
+
+    def init(self) -> dict[str, Any]:
+        """Return an empty projection — this transformer only suppresses events.
+
+        Returns:
+            An empty mapping; no `run.<key>` projection is added.
+        """
+        return {}
+
+    def process(self, event: ProtocolEvent) -> bool:
+        """Drop tagged `messages` events after nudging `MessagesTransformer` to ignore them.
+
+        Args:
+            event: The protocol event to observe, and possibly mutate.
+
+        Returns:
+            `False` to drop a tagged internal call from the raw event log and
+            `run.messages`; `True` otherwise.
+        """
+        if event["method"] != "messages":
+            return True
+
+        params = event["params"]
+        if params["namespace"] != self._scope_list:
+            return True
+
+        payload, metadata = params["data"]
+        is_internal = (
+            bool(metadata) and metadata.get(INTERNAL_CALL_METADATA_KEY) == _INTERNAL_CALL_TOKEN
+        )
+        if not is_internal:
+            return True
+
+        # Only `message-start` needs mutation: marking it as `"tool"` makes
+        # `MessagesTransformer` ignore the rest of that run. All events are still
+        # dropped from the raw log below.
+        if isinstance(payload, dict) and payload.get("event") == "message-start":
+            payload["role"] = "tool"
+        elif isinstance(payload, BaseMessage):
+            params["data"] = (None, metadata)
+
+        return False

coco/agents/middleware/model_call_limit.py

@@ -0,0 +1,267 @@
+"""Call tracking middleware for agents."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated, Any, Literal
+
+from coco_core.messages import AIMessage
+from coco_graph.channels.untracked_value import UntrackedValue
+from typing_extensions import NotRequired, override
+
+from coco.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ContextT,
+    PrivateStateAttr,
+    ResponseT,
+    hook_config,
+)
+
+if TYPE_CHECKING:
+    from coco_graph.runtime import Runtime
+
+
+class ModelCallLimitState(AgentState[ResponseT]):
+    """State schema for `ModelCallLimitMiddleware`.
+
+    Extends `AgentState` with model call tracking fields.
+
+    Type Parameters:
+        ResponseT: The type of the structured response. Defaults to `Any`.
+    """
+
+    thread_model_call_count: NotRequired[Annotated[int, PrivateStateAttr]]
+    run_model_call_count: NotRequired[Annotated[int, UntrackedValue, PrivateStateAttr]]
+
+
+def _build_limit_exceeded_message(
+    thread_count: int,
+    run_count: int,
+    thread_limit: int | None,
+    run_limit: int | None,
+) -> str:
+    """Build a message indicating which limits were exceeded.
+
+    Args:
+        thread_count: Current thread model call count.
+        run_count: Current run model call count.
+        thread_limit: Thread model call limit (if set).
+        run_limit: Run model call limit (if set).
+
+    Returns:
+        A formatted message describing which limits were exceeded.
+    """
+    exceeded_limits = []
+    if thread_limit is not None and thread_count >= thread_limit:
+        exceeded_limits.append(f"thread limit ({thread_count}/{thread_limit})")
+    if run_limit is not None and run_count >= run_limit:
+        exceeded_limits.append(f"run limit ({run_count}/{run_limit})")
+
+    return f"Model call limits exceeded: {', '.join(exceeded_limits)}"
+
+
+class ModelCallLimitExceededError(Exception):
+    """Exception raised when model call limits are exceeded.
+
+    This exception is raised when the configured exit behavior is `'error'` and either
+    the thread or run model call limit has been exceeded.
+    """
+
+    def __init__(
+        self,
+        thread_count: int,
+        run_count: int,
+        thread_limit: int | None,
+        run_limit: int | None,
+    ) -> None:
+        """Initialize the exception with call count information.
+
+        Args:
+            thread_count: Current thread model call count.
+            run_count: Current run model call count.
+            thread_limit: Thread model call limit (if set).
+            run_limit: Run model call limit (if set).
+        """
+        self.thread_count = thread_count
+        self.run_count = run_count
+        self.thread_limit = thread_limit
+        self.run_limit = run_limit
+
+        msg = _build_limit_exceeded_message(thread_count, run_count, thread_limit, run_limit)
+        super().__init__(msg)
+
+
+class ModelCallLimitMiddleware(
+    AgentMiddleware[ModelCallLimitState[ResponseT], ContextT, ResponseT]
+):
+    """Tracks model call counts and enforces limits.
+
+    This middleware monitors the number of model calls made during agent execution
+    and can terminate the agent when specified limits are reached. It supports
+    both thread-level and run-level call counting with configurable exit behaviors.
+
+    Thread-level: The middleware tracks the number of model calls and persists
+    call count across multiple runs (invocations) of the agent.
+
+    Run-level: The middleware tracks the number of model calls made during a single
+    run (invocation) of the agent.
+
+    Example:
+        ```python
+        from coco.agents.middleware import ModelCallLimitMiddleware
+        from coco.agents import create_agent
+
+        # Create middleware with limits
+        call_tracker = ModelCallLimitMiddleware(thread_limit=10, run_limit=5, exit_behavior="end")
+
+        agent = create_agent("openai:gpt-5.5", middleware=[call_tracker])
+
+        # Agent will automatically jump to end when limits are exceeded
+        result = await agent.invoke({"messages": [HumanMessage("Help me with a task")]})
+        ```
+    """
+
+    state_schema = ModelCallLimitState  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        *,
+        thread_limit: int | None = None,
+        run_limit: int | None = None,
+        exit_behavior: Literal["end", "error"] = "end",
+    ) -> None:
+        """Initialize the call tracking middleware.
+
+        Args:
+            thread_limit: Maximum number of model calls allowed per thread.
+
+                `None` means no limit.
+            run_limit: Maximum number of model calls allowed per run.
+
+                `None` means no limit.
+            exit_behavior: What to do when limits are exceeded.
+
+                - `'end'`: Jump to the end of the agent execution and
+                    inject an artificial AI message indicating that the limit was
+                    exceeded.
+                - `'error'`: Raise a `ModelCallLimitExceededError`
+
+        Raises:
+            ValueError: If both limits are `None` or if `exit_behavior` is invalid.
+        """
+        super().__init__()
+
+        if thread_limit is None and run_limit is None:
+            msg = "At least one limit must be specified (thread_limit or run_limit)"
+            raise ValueError(msg)
+
+        if exit_behavior not in {"end", "error"}:
+            msg = f"Invalid exit_behavior: {exit_behavior}. Must be 'end' or 'error'"
+            raise ValueError(msg)
+
+        self.thread_limit = thread_limit
+        self.run_limit = run_limit
+        self.exit_behavior = exit_behavior
+
+    @hook_config(can_jump_to=["end"])
+    @override
+    def before_model(
+        self, state: ModelCallLimitState[ResponseT], runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Check model call limits before making a model call.
+
+        Args:
+            state: The current agent state containing call counts.
+            runtime: The coco_graph runtime.
+
+        Returns:
+            If limits are exceeded and exit_behavior is `'end'`, returns
+                a `Command` to jump to the end with a limit exceeded message. Otherwise
+                returns `None`.
+
+        Raises:
+            ModelCallLimitExceededError: If limits are exceeded and `exit_behavior`
+                is `'error'`.
+        """
+        thread_count = state.get("thread_model_call_count", 0)
+        run_count = state.get("run_model_call_count", 0)
+
+        # Check if any limits will be exceeded after the next call
+        thread_limit_exceeded = self.thread_limit is not None and thread_count >= self.thread_limit
+        run_limit_exceeded = self.run_limit is not None and run_count >= self.run_limit
+
+        if thread_limit_exceeded or run_limit_exceeded:
+            if self.exit_behavior == "error":
+                raise ModelCallLimitExceededError(
+                    thread_count=thread_count,
+                    run_count=run_count,
+                    thread_limit=self.thread_limit,
+                    run_limit=self.run_limit,
+                )
+            if self.exit_behavior == "end":
+                # Create a message indicating the limit was exceeded
+                limit_message = _build_limit_exceeded_message(
+                    thread_count, run_count, self.thread_limit, self.run_limit
+                )
+                limit_ai_message = AIMessage(content=limit_message)
+
+                return {"jump_to": "end", "messages": [limit_ai_message]}
+
+        return None
+
+    @hook_config(can_jump_to=["end"])
+    async def abefore_model(
+        self,
+        state: ModelCallLimitState[ResponseT],
+        runtime: Runtime[ContextT],
+    ) -> dict[str, Any] | None:
+        """Async check model call limits before making a model call.
+
+        Args:
+            state: The current agent state containing call counts.
+            runtime: The coco_graph runtime.
+
+        Returns:
+            If limits are exceeded and exit_behavior is `'end'`, returns
+                a `Command` to jump to the end with a limit exceeded message. Otherwise
+                returns `None`.
+
+        Raises:
+            ModelCallLimitExceededError: If limits are exceeded and `exit_behavior`
+                is `'error'`.
+        """
+        return self.before_model(state, runtime)
+
+    @override
+    def after_model(
+        self, state: ModelCallLimitState[ResponseT], runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Increment model call counts after a model call.
+
+        Args:
+            state: The current agent state.
+            runtime: The coco_graph runtime.
+
+        Returns:
+            State updates with incremented call counts.
+        """
+        return {
+            "thread_model_call_count": state.get("thread_model_call_count", 0) + 1,
+            "run_model_call_count": state.get("run_model_call_count", 0) + 1,
+        }
+
+    async def aafter_model(
+        self,
+        state: ModelCallLimitState[ResponseT],
+        runtime: Runtime[ContextT],
+    ) -> dict[str, Any] | None:
+        """Async increment model call counts after a model call.
+
+        Args:
+            state: The current agent state.
+            runtime: The coco_graph runtime.
+
+        Returns:
+            State updates with incremented call counts.
+        """
+        return self.after_model(state, runtime)