使用 Upstash Ratelimit Callback 实现请求和令牌限速-CSDN博客

本文链接：https://blog.csdn.net/we4qwe1qwer4/article/details/144780883

老铁们，今天我来聊聊如何用 Upstash 的 RatelimitHandler 为你的应用添加请求或令牌限速。这个工具主要依赖 Upstash 的 ratelimit 库，而它的底层实现用到的是 Upstash Redis。

技术背景介绍

Upstash Ratelimit 的工作原理特别简单，其核心是每次调用 limit 方法时，都会向 Upstash Redis 发起一个 HTTP 请求。这时就会检查和更新用户的剩余请求或令牌数。如果剩余数不够了，我们可以阻止某些高耗费操作，比如调用大型语言模型（LLM）或查询向量存储：

response = ratelimit.limit()
if response.allowed:
    execute_costly_operation()

说白了，这就是个令牌桶算法，只是通过 Redis 做了分布式的实现。

原理深度解析

UpstashRatelimitHandler 可以轻松地将限速逻辑集成到你的工作链中。首先，你得去 Upstash 控制台创建一个 Redis 数据库（具体可以查看文档）。创建好数据库后，就要设置环境变量：

export UPSTASH_REDIS_REST_URL="****"
export UPSTASH_REDIS_REST_TOKEN="****"

然后需要安装相关库：

pip install upstash-ratelimit upstash-redis

实战代码演示

每请求限速

假设我们希望用户每分钟最多调用10次，可以这样实现：

import os
from langchain_community.callbacks import UpstashRatelimitError, UpstashRatelimitHandler
from langchain_core.runnables import RunnableLambda
from upstash_ratelimit import FixedWindow, Ratelimit
from upstash_redis import Redis

# 设置环境变量
os.environ["UPSTASH_REDIS_REST_URL"] = "****"
os.environ["UPSTASH_REDIS_REST_TOKEN"] = "****"

# 创建限速器
ratelimit = Ratelimit(
    redis=Redis.from_env(),
    limiter=FixedWindow(max_requests=10, window=60),
)

# 创建handler
user_id = "user_id"  # 获取用户ID的方法
handler = UpstashRatelimitHandler(identifier=user_id, request_ratelimit=ratelimit)

# 创建mock链
chain = RunnableLambda(str)

# 使用handler调用链
try:
    result = chain.invoke("Hello world!", config={"callbacks": [handler]})
except UpstashRatelimitError:
    print("Handling ratelimit.", UpstashRatelimitError)

在这波操作中，只有在请求允许的情况下高耗费操作才会被执行。

每令牌限速

如果你的工作链中包含 LLM，你也可以选择按令牌限速。它的原理是在调用 LLM 之前检查剩余令牌，如果不足则抛出 UpstashRatelimitError。调用后，根据返回的令牌使用信息更新剩余令牌数。

例子如下：

import os
from langchain_community.callbacks import UpstashRatelimitError, UpstashRatelimitHandler
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI
from upstash_ratelimit import FixedWindow, Ratelimit
from upstash_redis import Redis

# 设置环境变量
os.environ["UPSTASH_REDIS_REST_URL"] = "****"
os.environ["UPSTASH_REDIS_REST_TOKEN"] = "****"
os.environ["OPENAI_API_KEY"] = "****"

# 创建限速器
ratelimit = Ratelimit(
    redis=Redis.from_env(),
    limiter=FixedWindow(max_requests=500, window=60),
)

# 创建handler
user_id = "user_id"
handler = UpstashRatelimitHandler(identifier=user_id, token_ratelimit=ratelimit)

# 创建链
as_str = RunnableLambda(str)
model = ChatOpenAI()
chain = as_str | model

# 使用handler调用链
try:
    result = chain.invoke("Hello world!", config={"callbacks": [handler]})
except UpstashRatelimitError:
    print("Handling ratelimit.", UpstashRatelimitError)