ModelCallLimitMiddleware Class — langchain Architecture

Source Code

libs/langchain_v1/langchain/agents/middleware/model_call_limit.py lines 94–267

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 langchain.agents.middleware.call_tracking import ModelCallLimitMiddleware
        from langchain.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-4o", 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.