I am gonna use langchain builtin middleware in custom state graph

OSS Product Help LangGraph
,
1

is it possible to directly use the langchain builtin middleware like ModelCallLimitMiddleware in my custom stateGraph.

Much appreciate if you help me

2

hi @Dev

what do you mean by my custom stateGraph?

Afaik LangChain’s built‑in agent middleware (e.g., ModelCallLimitMiddleware) is wired into the agent graph that create_agent builds. In a hand-built StateGraph, there is no generic graph.add_middleware(...) and the built-ins assume the standard agent loop nodes (model, tools, end) and jump semantics.

I reckon you can instantiate ModelCallLimitMiddleware and call its hooks around your model node:

  • Call before_model(state, runtime) (or abefore_model) before the model; if it returns {"jump_to": "end", "messages": [...]}, route to your END node; if exit_behavior="error", handle ModelCallLimitExceededError
  • Call after_model(state, runtime) (or aafter_model) after the model to increment counters
  • Ensure your state carries thread_model_call_count and run_model_call_count; if you use a checkpointer, treat run count as non-persistent (untracked) yourself

Maybe sth like this:

from langchain.agents.middleware.model_call_limit import ModelCallLimitMiddleware, ModelCallLimitExceededError

limit = ModelCallLimitMiddleware(thread_limit=10, run_limit=5, exit_behavior="end")

def before_model_node(state, runtime):
    updates = limit.before_model(state, runtime) or {}
    # If updates includes {"jump_to": "end"}, add your conditional edge to END.
    return updates

def model_node(state):
    # ... call your model ...
    return {"messages": [...]}  # your AIMessage(s)

def after_model_node(state, runtime):
    return limit.after_model(state, runtime) or {}

Key caveats:

  • The middleware’s hook_config(can_jump_to=["end"]) metadata isn’t wired automatically in a plain StateGraph; you must add conditional edges based on the returned jump_to

  • The built-in uses UntrackedValue for the run counter; replicate that by not checkpointing run_model_call_count (or resetting it per run)

2 Likes
Implement ToolLimitationNode in stateGraph
3

Hello Pawel

I succeed to implement the langchain builtin middleware in the stateGraph

Thank

1 Like
4

@Dev Do you mind giving an example of how you are making this work?

5 
limit = ToolCallLimitMiddleware(
    thread_limit=None,
    run_limit=DEFAULT_TOOL_LIMIT,
    exit_behavior=DEFAULT_EXIT_BEHAVIOR,
)
def limit_check_node(state: GraphState, runtime):
        """
        MUST be placed after tool execution.
        This is the only correct place where tool call count is updated.
        """
        updates = limit.after_model(state, runtime) or {}

        if updates.get("jump_to") == "end":
            return {"__end__": True}

        return updates
1 Like
Implement ToolLimitationNode in stateGraph
6

Hello @pawel-twardziak
I have another significant challenge to implement the ToolCallLimitMiddleware
I set the tool call limit to 5. When I ask the LLM to create 6 images, the assistant returns:

tool_calls = [call1, call2, call3, call4, call5, call6]

What I want is to execute only the first 5 tool calls, ignore the 6th, and then return control to the LLM.

However, with my current implementation using ToolCallLimitMiddleware, all 6 tool calls are executed sequentially, and the middleware only triggers after the 6th call, which defeats the purpose of the limit.

How can I modify the implementation so that only the first 5 tool calls are executed and the 6th is skipped?

1 Like
7

Hi @Dev
Hmm that is an interesting issue :face_with_raised_eyebrow:
Could you create another post since this one is related to original issue and already resolved?
Just to keep things clear and well structured here on the forum. Appreciate :wink:
You can tag me there and I will answer today later on.

Implement ToolLimitationNode in stateGraph
8

Thank you, I just did it.

1 Like
9

I’m attaching the link here as the follow up