使用 Upstash 实现请求和令牌的限流

在现代应用程序中，限流是一种关键技术，可以防止系统过载并控制API的使用量。在这篇文章中，我们将学习如何使用 Upstash 的 RatelimitHandler 来实现基于请求数量或令牌数量的限流。这一工具借助 Upstash 提供的 Redis 实现，确保每次调用 limit 方法时，都能通过 HTTP 请求更新用户的剩余请求或令牌，并根据剩余数量来判断是否继续执行诸如调用大型语言模型（LLM）或者查询向量存储等昂贵操作。

技术背景介绍

Upstash 是一种托管的 Redis 云服务，提供了强大而简单的 API 接口，适合现代应用对 Redis 的使用场景。通过 Upstash 的 ratelimit 库，可以轻松实现请求或令牌的限流。

核心原理解析

限流的核心在于通过 Redis 的存储能力对用户请求进行计数。当用户的请求超过了预设的限额时，系统会停止用户的后续操作，从而避免系统负载过高。限流策略通常包括固定窗口、滑动窗口等，本文中我们将使用固定窗口策略（FixedWindow）。

代码实现演示

开始之前，请确保你已经注册并创建了 Upstash Redis 数据库，并设置好必要的环境变量：

export UPSTASH_REDIS_REST_URL='your-redis-url'
export UPSTASH_REDIS_REST_TOKEN='your-redis-token'

然后安装必要的 Python 库：

pip install upstash-ratelimit upstash-redis

每请求限流

假设我们希望每分钟允许用户最多进行 10 次请求，下面是实现该功能的代码示例：

import os

# 设置环境变量
os.environ["UPSTASH_REDIS_REST_URL"] = "your-redis-url"
os.environ["UPSTASH_REDIS_REST_TOKEN"] = "your-redis-token"

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

# 创建限流器
ratelimit = Ratelimit(
    redis=Redis.from_env(),
    # 每60秒最多10次请求
    limiter=FixedWindow(max_requests=10, window=60),
)

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

# 创建一个模拟链
chain = RunnableLambda(str)

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

每令牌限流

除了对请求进行限流，也可以根据令牌数量进行限流，这在使用 LLM 场景中特别有效。以下是实现动态令牌限流的代码：

import os

# 设置环境变量
os.environ["UPSTASH_REDIS_REST_URL"] = "your-redis-url"
os.environ["UPSTASH_REDIS_REST_TOKEN"] = "your-redis-token"
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"

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

# 创建限流器
ratelimit = Ratelimit(
    redis=Redis.from_env(),
    # 每60秒最多500个令牌
    limiter=FixedWindow(max_requests=500, window=60),
)

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

# 创建一个包含 LLM 的链
as_str = RunnableLambda(str)
model = ChatOpenAI()

chain = as_str | model

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

应用场景分析

这种限流机制特别适合像 API 网关、微服务架构中可能遇到的限流需求，并能有效防止恶意用户的滥用。

实践建议

确保你的 Redis 实例有足够的性能来处理限流请求。
选择适合你业务需求的限流策略，例如固定窗口或滑动窗口。
定期监控和调整限流参数，以适应不断变化的流量模式。

如果遇到问题欢迎在评论区交流。

—END—