Changing the call: wrap_model_call

wrap_model_call runs around each model call: it gets the request, can change it, and decides when to pass it on. dynamic_prompt uses it for the system prompt.

Lesson 16's hooks see the state. A wrap hook sees the request about to go to the model, with its messages, tools, system prompt and model, and a handler that sends it on. Calling handler(request) once is the normal path; not calling it, or calling it twice, is how middleware skips or retries a call.

Exampleprompts.py

from dataclasses import dataclass

from langchain.agents.middleware import dynamic_prompt, wrap_model_call


@dataclass
class Customer:
    name: str


@dynamic_prompt
def with_name(request):
    return f"You help {request.runtime.context.name}, a customer of a small online shop."


@wrap_model_call
def show_prompt(request, handler):
    print("system prompt:", request.system_prompt)
    return handler(request)

@dynamic_prompt turns a function into a wrap hook that sets the system prompt, here with the customer's name from the runtime context. show_prompt prints what the model is about to get and passes the request on unchanged.

Exampleagent.py

from langchain.agents import create_agent
from prompts import Customer, show_prompt, with_name
from shop_model import ShopModel
from tools import lookup_order

agent = create_agent(ShopModel(), tools=[lookup_order], context_schema=Customer,
                     middleware=[with_name, show_prompt])

Example

agent.invoke({"messages": [{"role": "user", "content": "Where is A17?"}]}, context=Customer("ravi"))

Two model calls, and each got a system prompt written for Ravi. with_name comes first in the list, so it is the outer layer: it set the prompt before show_prompt saw the request.

A different model per request

Examplegreeter.py

import re

from langchain.agents.middleware import wrap_model_call
from langchain.messages import AIMessage
from shop_model import ShopModel


class Greeter(ShopModel):
    def decide(self, messages):
        return AIMessage("Hi! Tell me an order number and I will look it up.")


@wrap_model_call
def pick_model(request, handler):
    if not re.search(r"\b[A-Z]\d+\b", request.messages[-1].text):
        request = request.override(model=Greeter())
    return handler(request)

request.override returns a copy of the request with one thing changed, here the model. A message with no order id goes to a model that only greets, which with hosted models is how a cheap model handles small talk and a larger one the real work.

Exampleagent.py

from langchain.agents import create_agent
from greeter import pick_model
from shop_model import ShopModel
from tools import lookup_order

agent = create_agent(ShopModel(), tools=[lookup_order], middleware=[pick_model])

Example

for text in ["Hello", "Where is A17?"]:
    result = agent.invoke({"messages": [{"role": "user", "content": text}]})
    print(result["messages"][-1].text)

Try it yourself

Change with_name to add the number of messages in request.messages to the prompt.
Make show_prompt skip the model by returning without calling handler, and read the error.
Swap the order of with_name and show_prompt and compare what is printed.

PreviousThe order middleware runs in

Slow is fine. Stopping is the only problem.