# 1. 安装依赖（Python 3.9+，指定 aiohttp 版本防兼容坑）
pip install locust openai python-dotenv aiohttp==3.9.1

# 2. 新建 .env 文件（存密钥，别硬编码！）
echo "OPENAI_API_KEY=sk-你的密钥
LLM_BASE_URL=https://api.你的模型地址/v1" > .env

# locust_llm.py
import os
import uuid
import time
import json
import tiktoken  # 精准计算 Token 的库
from dotenv import load_dotenv
from locust import HttpUser, task

# 加载环境变量
load_dotenv()
API_KEY = os.getenv("OPENAI_API_KEY")
BASE_URL = os.getenv("LLM_BASE_URL")

# 【新增】精准计算 Token 数（替代原"//2 估算"，避免误差）
def count_tokens(prompt, model="gpt-4-turbo"):
    encoder = tiktoken.encoding_for_model(model)
    return len(encoder.encode(prompt))

# 1. 生成测试用例（短/长 prompt 覆盖）
def generate_test_cases():
    # 场景1：电商客服（100 Token，高频）
    short_prompt = f"({uuid.uuid4()}) 我的订单12345没发货，帮我查物流"
    # 场景2：金融报告（16k Token，高资源消耗）
    long_prompt = f"({uuid.uuid4()}) " + "某公司2024财报：营收10亿，净利润2.3亿，毛利率35%..." * 80

    # 验证 Token 数（避免实际长度不符）
    print(f"短prompt Token数：{count_tokens(short_prompt)}")
    print(f"长prompt Token数：{count_tokens(long_prompt)}")

    return [
        {"name": "short_prompt", "content": short_prompt},
        {"name": "long_prompt", "content": long_prompt}
    ]

TEST_CASES = generate_test_cases()

# 2. 自定义指标容器（TTFT、吐字率，按 prompt 类型分）
CUSTOM_STATS = {
    "first_token_latency_short": [],
    "first_token_latency_long": [],
    "token_rate_short": [],
    "token_rate_long": []
}

# 3. 定义压测用户行为
class LLMTestUser(HttpUser):
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    # 权重 3:1，模拟电商场景短 prompt 更多
    @task(3)
    def test_short_prompt(self):
        self.run_test(TEST_CASES[0])

    @task(1)
    def test_long_prompt(self):
        self.run_test(TEST_CASES[1])

    def run_test(self, test_case):
        prompt_name = test_case["name"]
        prompt_content = test_case["content"]

        start_time = time.time()
        first_token_time = None
        total_tokens = 0

        # 流式请求（模拟用户"边看边等"的真实场景）
        with self.client.post(
            url=f"{BASE_URL}/chat/completions",
            json={
                "model": "gpt-4-turbo",  # 替换成你的模型（如 qwen-plus/vllm-7b）
                "messages": [{"role": "user", "content": prompt_content}],
                "stream": True,
                "max_tokens": 1000  # 按业务需求改
            },
            headers=self.headers,
            stream=True,
            name=prompt_name  # Locust UI 里按 prompt 类型分组
        ) as response:
            for line in response.iter_lines():
                if not line:
                    continue
                line_str = line.decode("utf-8").strip()
                if not line_str.startswith("data:"):
                    continue
                line_str = line_str[5:].strip()
                if line_str == "[DONE]":
                    break
                # 解析 Token（按你模型的响应格式调整，此处适配 OpenAI）
                try:
                    data = json.loads(line_str)
                    token = data["choices"][0]["delta"].get("content", "")
                    if token and not first_token_time:
                        # 记录首 Token 到达时间
                        first_token_time = time.time()
                    if token:
                        # 用 tiktoken 精准算 Token 数（替代原估算）
                        total_tokens += count_tokens(token)
                except Exception as e:
                    print(f"解析响应出错：{e}")

        # 计算并记录指标
        if first_token_time and total_tokens > 0:
            first_token_latency = first_token_time - start_time
            token_rate = total_tokens / (time.time() - first_token_time)

            # 按 prompt 类型存指标
            if prompt_name == "short_prompt":
                CUSTOM_STATS["first_token_latency_short"].append(first_token_latency)
                CUSTOM_STATS["token_rate_short"].append(token_rate)
            else:
                CUSTOM_STATS["first_token_latency_long"].append(first_token_latency)
                CUSTOM_STATS["token_rate_long"].append(token_rate)

# 启动命令（复制到终端执行）
# locust -f locust_llm.py --host=https://api.你的模型地址 --web-port=8089

stages:
  - duration: 3m  # 预热期：1 并发（让模型加载缓存）
    target: 1
  - duration: 10m # 基准期：10 并发（日常流量）
    target: 10
  - duration: 10m # 压力期：25 并发（峰值流量）
    target: 25
  - duration: 5m  # 极限期：40 并发（故障演练）
    target: 40