ToolCallLimitMiddleware

toolcalllimitmiddleware.py

from langchain.agents.middleware import AgentMiddleware, AgentState, ModelRequest

class ToolCallLimitMiddleware(AgentMiddleware):
    def __init__(self, limit) -> None:
        super().__init__()
        self.limit = limit

    def modify_model_request(
        self, request: ModelRequest, state: AgentState
    ) -> ModelRequest:
        tool_call_count = sum(
            1
            for msg in state["messages"]
            if isinstance(msg, AIMessage) and msg.tool_calls
        )
        if tool_call_count >= self.limit:
            logger.info(
                "Tool call limit of %d reached, disabling further tool calls.",
                self.limit,
            )
            request.tools = []
            request.messages.append(
                HumanMessage(
                    content="No more tool calls can be made, now it's time to summarize your findings based off previous research."
                )
            )
        return request

toolcalllimitmiddleware_diagram.mermaid