Skip to content

AI budget control for LangChain

AI providers bill per token. Without per-user limits, a single user can exhaust your entire monthly budget – through prompt attacks, runaway loops, or just heavy legitimate use.

A token bucket rate limit maps directly onto how AI billing works. You estimate the cost of each request in tokens, deduct it from the user’s bucket, and deny requests when the bucket is empty. The bucket refills over time, giving each user a sustained allowance without sharp rate-limit cliffs.

Alternatively, you can use a fixed window or sliding window limit to enforce a hard cap on spend per user per day, week, or month. For details on different approaches, see the rate limiting algorithms reference.

Arcjet handles bucket state across all instances of your application – no Redis or external state store required.

This example deducts an estimated token cost from a per-user bucket before a tool runs. When the bucket is empty, the tool does not execute.

We assume you already have a LangChain project set up. For helper options and denial behavior, see the LangChain agent guard.

Install the dependencies:

Terminal window
# Export your Arcjet API key from https://console.arcjet.com
export ARCJET_KEY="ajkey_..."
export ARCJET_ENV=development
pip install "arcjet[langchain-agents]" langchain langchain-openai

Create the example:

agent.py
import os
from arcjet.guard import TokenBucket, launch_arcjet
from arcjet.guard.langchain import guard_tool
from langchain_core.tools import tool
arcjet = launch_arcjet(key=os.environ["ARCJET_KEY"])
token_budget = TokenBucket(
refill_rate=2000,
interval_seconds=3600,
max_tokens=5000,
bucket="ai-tokens",
)
@tool
async def complete_prompt(prompt: str, estimated_tokens: int) -> dict:
"""Complete a user prompt."""
return {"prompt": prompt}
complete_prompt = guard_tool(
guard=arcjet,
tool=complete_prompt,
action="prompt.completed",
rules=lambda arguments, _config: [
token_budget(
key="user123", # Replace with your authenticated user ID
requested=max(1, int(arguments["estimated_tokens"])),
)
],
)

Then start or invoke the agent with a test prompt.

Requests appear in your Arcjet dashboard in real time.

characteristics: ["userId"] - Tracks the bucket per user. Replace "userId" with the characteristic that identifies a unique user in your application, such as a session token, API key, or authenticated user ID. Pass the value to aj.protect() as a named argument.

refillRate and interval - Set the sustained allowance. refillRate: 2_000, interval: "1h" gives each user 2,000 tokens per hour. Adjust to match your AI provider’s pricing and your cost targets. These are hard coded in this example, but you can also calculate them dynamically based on user subscription level or other factors. Pass the calculated values to the rule.

capacity - The maximum tokens a user can accumulate. Setting capacity: 5_000 with refillRate: 2_000 lets users burst up to 5,000 tokens if they haven’t used their allowance recently.

The example uses a characters / 4 heuristic (~1 token per 4 characters for common English text). This is a reasonable starting point – it avoids introducing extra dependencies and works well enough for budget enforcement where a small margin of error is acceptable.

For accurate counts, use a tokenizer: