Middleware & Guardrails

from standins import ScriptedModel, RuleModel, call
from langchain.agents import create_agent
from langchain.tools import tool
from langchain.messages import AIMessage
from langchain.agents.middleware import (PIIMiddleware, HumanInTheLoopMiddleware, ToolCallLimitMiddleware,
                                         before_model, after_model, wrap_tool_call, AgentState)
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command
from langgraph.runtime import Runtime

@tool
def get_balance(account_id: str) -> str:
    """Get account balance."""
    return f"Account {account_id} balance: Rs 1,24,500"

@tool
def transfer_money(to_account: str, amount: float) -> str:
    """Transfer money to another account."""
    return f"Rs {amount} transferred to {to_account}."
bank_tools = [get_balance, transfer_money]

print("--- 1. PII middleware ---")
class Echo(RuleModel):
    def _pick(self, messages): return "You shared: " + messages[-1].content
pii_agent = create_agent(Echo(), tools=[], middleware=[
    PIIMiddleware("email", strategy="redact", apply_to_input=True),
    PIIMiddleware("credit_card", strategy="mask", apply_to_input=True)])
r = pii_agent.invoke({"messages": [{"role": "user", "content": "My email is ajeet@example.com and card 4111 1111 1111 1111."}]})
print("what the LLM saw:", r["messages"][0].content)

print("\n--- 2. Human-in-the-loop ---")
for decision in ["approve", "reject"]:
    llm = ScriptedModel([call("transfer_money", {"to_account": "ACC999", "amount": 5000}),
                         "Done: Rs 5000 sent to ACC999." if decision == "approve" else "The transfer was not done: the manager rejected it."])
    hitl_agent = create_agent(llm, bank_tools, checkpointer=InMemorySaver(),
                              middleware=[HumanInTheLoopMiddleware(interrupt_on={"transfer_money": True, "get_balance": False})])
    config = {"configurable": {"thread_id": "hitl-" + decision}}
    r = hitl_agent.invoke({"messages": [{"role": "user", "content": "Transfer Rs 5000 to ACC999"}]}, config)
    req = r["__interrupt__"][0].value["action_requests"][0]
    print(f"⏸ paused, waiting for approval: {req['name']}({req['args']})")
    resume = {"decisions": [{"type": "approve"}]} if decision == "approve" else {"decisions": [{"type": "reject", "message": "Manager rejected this transfer."}]}
    r = hitl_agent.invoke(Command(resume=resume), config)
    print(f"  {decision:7s} -> tool said: {r['messages'][-2].content!r}\n            agent said: {r['messages'][-1].content}")

print("\n--- 3. Tool call limit ---")
greedy = ScriptedModel([call("transfer_money", {"to_account": "A1", "amount": 100}, "t1"), call("transfer_money", {"to_account": "A2", "amount": 100}, "t2"), "All done."])
r = create_agent(greedy, bank_tools, middleware=[ToolCallLimitMiddleware(tool_name="transfer_money", run_limit=1)]).invoke(
    {"messages": [{"role": "user", "content": "Send 100 to A1 and 100 to A2"}]})
print([m.content[:60] for m in r["messages"] if m.type == "tool"])

print("\n--- 4. Custom logging + compliance guard ---")
@before_model
def log_before(state: AgentState, runtime: Runtime):
    print(f"➡️  calling LLM with {len(state['messages'])} messages")

@wrap_tool_call
def log_tools(request, handler):
    print(f"🔧 tool: {request.tool_call['name']} {request.tool_call['args']}")
    return handler(request)

log_llm = ScriptedModel([call("get_balance", {"account_id": "ACC777"}), "Balance of ACC777 is Rs 1,24,500."])
print(create_agent(log_llm, bank_tools, middleware=[log_before, log_tools]).invoke({"messages": [{"role": "user", "content": "Balance of ACC777?"}]})["messages"][-1].content)

BANNED = ["guaranteed returns", "100% profit"]
@after_model(can_jump_to=["end"])
def compliance_guard(state: AgentState, runtime: Runtime):
    text = (state["messages"][-1].content or "").lower()
    if any(word in text for word in BANNED):
        return {"messages": [AIMessage("⚠️ This answer was blocked by compliance policy.")], "jump_to": "end"}
pushy = ScriptedModel(["Invest now! Our fund gives guaranteed returns of 30%!"])
print(create_agent(pushy, [], middleware=[compliance_guard]).invoke({"messages": [{"role": "user", "content": "Should I invest in your mutual fund?"}]})["messages"][-1].content)