Middleware

from typing import Any

from langchain.agents.middleware import AgentState, after_model, before_model
from langchain.messages import AIMessage
from langgraph.runtime import Runtime
from typing_extensions import NotRequired


class CounterState(AgentState):
    model_calls: NotRequired[int]


@before_model(state_schema=CounterState, can_jump_to=["end"])
def stop_after_two(state: CounterState, runtime: Runtime) -> dict[str, Any] | None:
    if state.get("model_calls", 0) >= 2:
        return {"messages": [AIMessage("Stopping: model-call budget used.")], "jump_to": "end"}
    return None


@after_model(state_schema=CounterState)
def count_calls(state: CounterState, runtime: Runtime) -> dict[str, Any] | None:
    return {"model_calls": state.get("model_calls", 0) + 1}

from typing import Callable

from langchain.agents import AgentState, create_agent
from langchain.agents.middleware import ExtendedModelResponse, ModelRequest, ModelResponse, wrap_model_call
from langgraph.types import Command
from typing_extensions import NotRequired


class UsageState(AgentState):
    last_answer_chars: NotRequired[int]


@wrap_model_call(state_schema=UsageState)
def track_answer_size(request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ExtendedModelResponse:
    response = handler(request)
    size = len(response.result[-1].content)
    return ExtendedModelResponse(model_response=response, command=Command(update={"last_answer_chars": size}))


agent = create_agent(model, tools=[], middleware=[track_answer_size])
out = agent.invoke({"messages": [{"role": "user", "content": "What is HNSW?"}]})
print(out["last_answer_chars"])

from langchain.agents import create_agent
from langchain.agents.middleware import SummarizationMiddleware
from langgraph.checkpoint.memory import InMemorySaver

agent = create_agent(
    answerer,                                  # scripted fake model: "answer 1", "answer 2", ...
    tools=[],
    checkpointer=InMemorySaver(),
    middleware=[SummarizationMiddleware(model=summarizer, trigger=("messages", 4), keep=("messages", 2))],
)
cfg = {"configurable": {"thread_id": "t1"}}
for q in ["What is HNSW?", "And ef_search?", "And M?"]:
    out = agent.invoke({"messages": [{"role": "user", "content": q}]}, cfg)
    print(len(out["messages"]), [f"{type(m).__name__}:{str(m.content)[:40]}" for m in out["messages"]])

2 ['HumanMessage:What is HNSW?', 'AIMessage:answer 1']
4 ['HumanMessage:What is HNSW?', 'AIMessage:answer 1', 'HumanMessage:And ef_search?', 'AIMessage:answer 2']
4 ['HumanMessage:Here is a summary of the conversation to', 'AIMessage:answer 2', 'HumanMessage:And M?', 'AIMessage:answer 3']

from langchain.agents import create_agent
from langchain.agents.middleware import HumanInTheLoopMiddleware
from langchain.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command

opened = []


@tool
def open_github_issue(title: str, body: str) -> str:
    """Open a GitHub issue in the docs repository."""
    opened.append(title)
    return f"Opened issue: {title}"


agent = create_agent(
    model,                       # scripted fake model: one open_github_issue call, then "Issue opened."
    tools=[open_github_issue],
    checkpointer=InMemorySaver(),
    middleware=[HumanInTheLoopMiddleware(
        interrupt_on={"open_github_issue": {"allowed_decisions": ["approve", "edit", "reject"]}},
        description_prefix="DocsAssistant wants to open an issue",
    )],
)
cfg = {"configurable": {"thread_id": "t1"}}
result = agent.invoke({"messages": [{"role": "user", "content": "File a docs issue about ef_search"}]}, cfg, version="v2")
req = result.interrupts[0].value
print("action_requests:", req["action_requests"])
print("review_configs:", req["review_configs"])
print("issues opened before approval:", opened)
out = agent.invoke(Command(resume={"decisions": [{"type": "approve"}]}), cfg, version="v2")
print("issues opened after approval:", opened)
print(out.value["messages"][-1].content)

action_requests: [{'name': 'open_github_issue', 'args': {'title': 'Doc gap: ef_search', 'body': 'Explain ef_search.'}, 'description': "DocsAssistant wants to open an issue\n\nTool: open_github_issue\nArgs: {'title': 'Doc gap: ef_search', 'body': 'Explain ef_search.'}"}]
review_configs: [{'action_name': 'open_github_issue', 'allowed_decisions': ['approve', 'edit', 'reject']}]
issues opened before approval: []
issues opened after approval: ['Doc gap: ef_search']
Issue opened.

from langchain.agents import create_agent
from langchain.agents.middleware import PIIDetectionError, PIIMiddleware

agent = create_agent(model, tools=[], middleware=[
    PIIMiddleware("email", strategy="redact"),
    PIIMiddleware("credit_card", strategy="mask"),
])
agent.invoke({"messages": [{"role": "user", "content": "Mail ana@example.com, card 4111 1111 1111 1111"}]})
print("model saw:", model.seen[0][-1].content)      # `model` here records what it was sent

strict = create_agent(model2, tools=[], middleware=[
    PIIMiddleware("api_key", detector=r"sk-[a-zA-Z0-9]{32}", strategy="block"),
])
try:
    strict.invoke({"messages": [{"role": "user", "content": user_text_with_a_key}]})
except PIIDetectionError as e:
    print("blocked:", type(e).__name__, "-", e)

model saw: Mail [REDACTED_EMAIL], card **** **** **** 1111
blocked: PIIDetectionError - Detected 1 instance(s) of api_key in text content

from langchain.agents import create_agent
from langchain.agents.middleware import ModelRetryMiddleware, ToolRetryMiddleware

agent = create_agent(
    flaky_model,                 # fake model: raises ConnectionError twice, then works
    tools=[search_docs],         # raises TimeoutError twice, then returns
    middleware=[
        ModelRetryMiddleware(max_retries=3, initial_delay=0.0, jitter=False),   # delays disabled for the demo
        ToolRetryMiddleware(max_retries=3, initial_delay=0.0, jitter=False),
    ],
)
out = agent.invoke({"messages": [{"role": "user", "content": "What is HNSW?"}]})
print("model failures left:", flaky_model.failures, "| tool attempts:", attempts["n"], "|", out["messages"][-1].content)

# retries exhausted: on_failure="continue" (default) turns the failure into an AIMessage
agent = create_agent(always_failing_model, tools=[], middleware=[ModelRetryMiddleware(max_retries=1, initial_delay=0.0, jitter=False)])
out = agent.invoke({"messages": [{"role": "user", "content": "hi"}]})
print("exhausted:", out["messages"][-1].content)

model failures left: 0 | tool attempts: 3 | Done.
exhausted: Model call failed after 2 attempts with ConnectionError: ollama: connection reset

from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware, ToolCallLimitMiddleware
from langgraph.checkpoint.memory import InMemorySaver

agent = create_agent(
    looping_model,               # fake model that requests get_page forever
    tools=[get_page], checkpointer=InMemorySaver(),
    middleware=[ModelCallLimitMiddleware(run_limit=3, exit_behavior="end")],
)
out = agent.invoke({"messages": [{"role": "user", "content": "loop"}]}, {"configurable": {"thread_id": "a"}})
print(out["messages"][-1].content)

agent = create_agent(
    three_calls_model,           # fake model: three get_page calls, then "Enough."
    tools=[get_page],
    middleware=[ToolCallLimitMiddleware(tool_name="get_page", run_limit=2)],
)
out = agent.invoke({"messages": [{"role": "user", "content": "loop"}]})
for m in out["messages"]:
    if m.type == "tool":
        print(m.status, "|", m.content)

Model call limits exceeded: run limit (3/3)
success | <p0>
success | <p1>
error | Tool call limit exceeded. Do not call 'get_page' again.

from dataclasses import dataclass

from langchain.agents.middleware import ModelRequest, dynamic_prompt


@dataclass
class DocsContext:
    component: str
    expertise: str


@dynamic_prompt
def docs_prompt(request: ModelRequest) -> str:
    ctx = request.runtime.context
    base = f"You are DocsAssistant. Answer only from the {ctx.component} documentation and cite page paths."
    if ctx.expertise == "beginner":
        base += " Explain terms briefly."
    return base

from typing import Callable

from langchain.agents.middleware import ModelRequest, ModelResponse, wrap_model_call

MAX_MESSAGES = 6


@wrap_model_call
def trim_history(request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelResponse:
    """Send the model only the most recent messages; the saved state keeps everything."""
    if len(request.messages) > MAX_MESSAGES:
        request = request.override(messages=request.messages[-MAX_MESSAGES:])
    return handler(request)

from typing import Callable

from langchain.agents import create_agent
from langchain.agents.middleware import ModelRequest, ModelResponse, ToolCallRequest, wrap_model_call, wrap_tool_call
from langchain.messages import ToolMessage
from langgraph.types import Command


@wrap_model_call
def hide_get_page_until_search(request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelResponse:
    searched = any(isinstance(m, ToolMessage) and m.name == "search_docs" for m in request.messages)
    if not searched:
        request = request.override(tools=[t for t in request.tools if t.name != "get_page"])
    return handler(request)


@wrap_tool_call
def audit_tool_calls(request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage | Command]) -> ToolMessage | Command:
    print("tool call:", request.tool_call["name"], request.tool_call["args"])
    result = handler(request)
    print("tool result:", getattr(result, "content", result))
    return result


agent = create_agent(
    model,                       # scripted fake: search_docs, then get_page, then a third call that is cut off
    tools=[search_docs, get_page],
    context_schema=DocsContext,
    middleware=[docs_prompt, stop_after_two, count_calls, hide_get_page_until_search, audit_tool_calls],
)
out = agent.invoke(
    {"messages": [{"role": "user", "content": "What is ef_search?"}]},
    context=DocsContext(component="vector-rag", expertise="beginner"),
)
print("system prompt:", model.seen[0][0].content)
print("tools bound per model call:", model.tools_seen)
print("model_calls:", out["model_calls"])

tool call: search_docs {'query': 'HNSW'}
tool result: 3 sections match 'HNSW'
tool call: get_page {'path': 'modules/ROOT/pages/database/vector-rag/ann-indexes.adoc'}
tool result: <full text of modules/ROOT/pages/database/vector-rag/ann-indexes.adoc>
system prompt: You are DocsAssistant. Answer only from the vector-rag documentation and cite page paths. Explain terms briefly.
tools bound per model call: [['search_docs'], ['search_docs', 'get_page']]
model_calls: 2