💰 设置团队预算

"""
- 运行两个并发团队调用相同模型
- 模型有60 TPM（每分钟令牌数）
- 模拟响应返回每次请求30个总令牌
- 每个团队每分钟只能发出一次请求
"""

import requests
from openai import OpenAI, RateLimitError

def create_key(api_key: str, base_url: str): 
    response = requests.post(
        url="{}/key/generate".format(base_url), 
        json={},
        headers={
            "Authorization": "Bearer {}".format(api_key)
        }
    )

    _response = response.json()

    return _response["key"]

key_1 = create_key(api_key="sk-1234", base_url="http://0.0.0.0:4000")
key_2 = create_key(api_key="sk-1234", base_url="http://0.0.0.0:4000")

# 使用key 1调用代理 - 成功
openai_client_1 = OpenAI(api_key=key_1, base_url="http://0.0.0.0:4000")

response = openai_client_1.chat.completions.with_raw_response.create(
    model="my-fake-model", messages=[{"role": "user", "content": "Hello world!"}],
)

print("调用1的响应头 - {}".format(response.headers))
_response = response.parse()
print("调用1的总令牌数 - {}".format(_response.usage.total_tokens))


# 使用key 2调用代理 - 成功
openai_client_2 = OpenAI(api_key=key_2, base_url="http://0.0.0.0:4000")

response = openai_client_2.chat.completions.with_raw_response.create(
    model="my-fake-model", messages=[{"role": "user", "content": "Hello world!"}],
)

print("调用2的响应头 - {}".format(response.headers))
_response = response.parse()
print("调用2的总令牌数 - {}".format(_response.usage.total_tokens))
# 使用key 2调用代理 - 失败
try:  
    openai_client_2.chat.completions.with_raw_response.create(model="my-fake-model", messages=[{"role": "user", "content": "Hey, how's it going?"}])
    raise Exception("这应该失败！")
except RateLimitError as e: 
    print("这被限流了，因为 - {}".format(str(e)))