demo_model_call_limit_middleware.py

from langchain.tools import ToolRuntime
from typing_extensions import TypedDict
from langchain.agents import create_agent
from langchain.agents.middleware.model_call_limit import (
    ModelCallLimitMiddleware,
    ModelCallLimitExceededError,
)
from langchain_core.tools import tool


class UserContext(TypedDict):
    user_id: str


# Mock database of user accounts
_MOCK_ACCOUNTS = {
    "user_123": {
        "name": "Alice Johnson",
        "email": "alice@example.com",
        "balance": 150.00,
        "plan": "Pro",
        "billing_date": "2025-01-15",
        "last_payment": "2024-12-15",
        "payment_method": "Credit card ending in 4242",
        "subscription_status": "active",
        "recent_charges": [
            {"date": "2024-12-15", "amount": 29.99, "description": "Pro Plan Monthly"},
            {"date": "2024-11-15", "amount": 29.99, "description": "Pro Plan Monthly"},
            {"date": "2024-10-15", "amount": 29.99, "description": "Pro Plan Monthly"},
        ],
    }
}

@tool
def query_knowledge_base() -> str:
    """Return knowledge base information."""
    return """
    * duplicate charges often occur due to payment processing retries. Check your bank statement to verify actual charges.
    * to cancel your subscription, go to Account Settings > Billing > Cancel Subscription. You'll retain access until your billing date.
    * we offer 30-day refunds for Pro plan upgrades. Refunds are processed within 5-10 business days.
    * update your billing address in Account Settings > Billing > Payment Method.
    """


@tool
def get_user_account(runtime: ToolRuntime[UserContext]) -> str:
    """Retrieve user account information including billing history."""
    user_id = runtime.context["user_id"]
    if user_id not in _MOCK_ACCOUNTS:
        return f"Account not found for user ID: {user_id}"
    account = _MOCK_ACCOUNTS[user_id]
    return (
        f"Account: {account['name']} | Email: {account['email']} | "
        f"Plan: {account['plan']} | Status: {account['subscription_status']} | "
        f"Balance: ${account['balance']:.2f} | "
        f"Last Payment: {account['last_payment']} | "
        f"Billing Date: {account['billing_date']}"
    )


_billing_subagent = create_agent(
    model="openai:gpt-4o-mini",
    tools=[get_user_account, query_knowledge_base],
    middleware=[
        ModelCallLimitMiddleware(
            run_limit=2,
            exit_behavior="error",
        )
    ],
    system_prompt="""You are a billing support agent. Help customers with their billing questions using the available tools.
    Please do not call multiple tools at once.""",
    context_schema=UserContext,
)


@tool
def escalate_to_human(issue_summary: str) -> str:
    """Escalate the query to a human support agent."""
    print(f"\n  🚨 ESCALATED TO HUMAN AGENT 🚨")
    print(f"  Issue: {issue_summary}")
    return (
        f"Your issue has been escalated to our support team. A specialist will assist you shortly."
    )


@tool
def billing_subagent(query: str, runtime: ToolRuntime) -> str:
    """
    Invokes the billing subagent to handle complex queries.
    If it hits its model call limit, we catch the error and escalate.
    """
    print(f"\n  → Delegating to billing subagent: '{query}'")

    try:
        result = _billing_subagent.invoke(
            {"messages": [{"role": "user", "content": query}]}, context=runtime.context
        )

        # Extract the final response
        final_message = result["messages"][-1]
        return f"Subagent result: {final_message.content}"

    except ModelCallLimitExceededError as e:
        return "The billing subagent cannot handle this query. Please escalate to a human agent."


main_agent = create_agent(
    model="openai:gpt-4o-mini",
    tools=[billing_subagent, escalate_to_human],
    system_prompt=(
        "You are a customer support agent. Use the billing_subagent tool "
        "to help customers with their billing questions. If the subagent indicates "
        "escalation is needed, use the escalate_to_human tool."
        "Please don't call multiple tools at once, but you can invoke the billing subagent with a complex query."
    ),
    context_schema=UserContext,
)

langgraph.json

{
  "dependencies": ["."],
  "graphs": {
    "model_call_limit": "./demo_model_call_limit.py:main_agent",
  },
  "env": ".env"
}