性能与压力测试（Performance & Stress Testing）

# 基于 openai==1.30+ 验证（截至 2026-03）
import time
import statistics
from dataclasses import dataclass, field
from openai import OpenAI

@dataclass
class AgentMetrics:
    """单次请求的性能指标"""
    ttft: float = 0.0           # 首Token延迟（秒）
    total_time: float = 0.0     # 端到端延迟（秒）
    input_tokens: int = 0       # 输入Token数
    output_tokens: int = 0      # 输出Token数
    tpot: float = 0.0           # 单Token生成时间（秒）

def measure_single_request(client: OpenAI, prompt: str, model: str = "gpt-4o") -> AgentMetrics:
    """
    测量单次LLM调用的性能指标。
    使用流式响应来精确测量TTFT和TPOT。
    """
    metrics = AgentMetrics()
    start = time.perf_counter()
    first_token_received = False
    token_count = 0

    # 流式调用，逐Token计时
    stream = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        stream=True,
        stream_options={"include_usage": True},
    )

    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            if not first_token_received:
                metrics.ttft = time.perf_counter() - start
                first_token_received = True
            token_count += 1

        # 从最后一个chunk获取usage信息
        if chunk.usage:
            metrics.input_tokens = chunk.usage.prompt_tokens
            metrics.output_tokens = chunk.usage.completion_tokens

    metrics.total_time = time.perf_counter() - start

    # 计算TPOT：(总时间 - TTFT) / 输出Token数
    if token_count > 1:
        metrics.tpot = (metrics.total_time - metrics.ttft) / (token_count - 1)

    return metrics

def run_baseline_test(client: OpenAI, prompts: list[str], model: str = "gpt-4o") -> dict:
    """
    运行基准测试，返回P50/P95/P99等汇总指标。
    prompts: 真实场景的Prompt列表（禁止使用随机文本）。
    """
    all_metrics = [measure_single_request(client, p, model) for p in prompts]

    ttfts = [m.ttft for m in all_metrics]
    totals = [m.total_time for m in all_metrics]
    tokens = [m.input_tokens + m.output_tokens for m in all_metrics]
    costs = [m.input_tokens * 2.50 / 1_000_000 + m.output_tokens * 10.00 / 1_000_000
             for m in all_metrics]  # GPT-4o 价格（美元）

    def percentile(data, p):
        sorted_data = sorted(data)
        idx = int(len(sorted_data) * p / 100)
        return sorted_data[min(idx, len(sorted_data) - 1)]

    return {
        "requests": len(all_metrics),
        "ttft_p50": f"{percentile(ttfts, 50)*1000:.0f}ms",
        "ttft_p95": f"{percentile(ttfts, 95)*1000:.0f}ms",
        "e2e_p50": f"{percentile(totals, 50):.2f}s",
        "e2e_p95": f"{percentile(totals, 95):.2f}s",
        "e2e_p99": f"{percentile(totals, 99):.2f}s",
        "avg_tokens": f"{statistics.mean(tokens):.0f}",
        "avg_cost": f"${statistics.mean(costs):.4f}",
        "total_cost": f"${sum(costs):.4f}",
    }

# 基于 locust==2.31+ 验证（截至 2026-03）
# 文件名: locustfile.py
from locust import HttpUser, task, between
import random

# 真实场景Prompt池（禁止使用随机文本）
PROMPT_POOL = [
    "帮我总结一下这篇文章的要点",           # 短Prompt
    "根据以下需求文档，设计一个API接口方案：" + "x" * 500,  # 长Prompt
    "对比 Redis 和 PostgreSQL 做向量存储的优劣",           # 中等Prompt
]

class AgentLoadUser(HttpUser):
    """模拟Agent用户的负载测试"""
    wait_time = between(1, 5)  # 模拟用户思考间隔

    @task(3)
    def single_turn(self):
        """单轮查询（权重3）"""
        self.client.post("/v1/chat/completions", json={
            "model": "gpt-4o",
            "messages": [{"role": "user", "content": random.choice(PROMPT_POOL)}],
        }, name="single_turn")

    @task(1)
    def multi_turn(self):
        """多轮对话（权重1），测试上下文累积效应"""
        messages = []
        for i in range(3):
            messages.append({"role": "user", "content": f"第{i+1}轮: {random.choice(PROMPT_POOL)}"})
            self.client.post("/v1/chat/completions", json={
                "model": "gpt-4o",
                "messages": messages.copy(),
            }, name=f"multi_turn_round_{i+1}")
            messages.append({"role": "assistant", "content": "模拟回复..."})

# 运行命令（渐进式加压）:
# locust -f locustfile.py --host=http://localhost:8000 \
#   --users 100 --spawn-rate 10 --run-time 10m
# --spawn-rate 10 表示每秒增加10个用户，10秒后达到100并发