llmspeed.py

import csv
import time
import yaml
import sys
from datetime import datetime
from typing import Optional

try:
    from openai import OpenAI
except ImportError:
    OpenAI = None

try:
    import anthropic
except ImportError:
    anthropic = None


def test_openai(client: OpenAI, model: str, prompt: str) -> dict:
    """测试 OpenAI Chat 格式模型 / Test OpenAI Chat format models"""
    start_time = time.time()
    first_token_time = None
    total_tokens = 0
    
    # o1 系列模型不支持 stream 参数
    # o1 series models don't support stream parameter
    if model.startswith('o1'):
        try:
            response = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
            )
            content = response.choices[0].message.content or ""
            first_token_time = time.time() - start_time
            total_tokens = estimate_tokens(content)
            total_time = time.time() - start_time
            return {
                "first_token_latency": round(first_token_time * 1000, 1),
                "tokens_per_second": round(total_tokens / total_time, 2) if total_time > 0 else 0,
                "total_tokens": int(total_tokens),
                "total_time": round(total_time, 3)
            }
        except Exception as e:
            # o1-mini 等可能支持 / o1-mini etc. may support stream
            pass
    
    # 标准流式响应 / Standard streaming response
    try:
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            stream=True,
        )
        
        # 检查响应是否为 SSE 流，防止代理返回 HTML 错误页
        # Check if response is SSE stream, prevent proxy returning HTML error page
        if hasattr(response, 'response'):
            raw = response.response
            ct = raw.headers.get('content-type', '')
            if 'text/event-stream' not in ct:
                print(f"      模型 {model} 不可用 / Model {model} unavailable (content-type: {ct})")
                return {
                    "first_token_latency": -1,
                    "tokens_per_second": -1,
                    "total_tokens": 0,
                    "total_time": round(time.time() - start_time, 3)
                }
        
        for chunk in response:
            if not hasattr(chunk, 'choices') or not chunk.choices:
                continue
            
            delta = chunk.choices[0].delta if chunk.choices else None
            if delta is None:
                continue
            
            content = getattr(delta, 'content', None)
            reasoning = getattr(delta, 'reasoning_content', None)
            effective_content = content or reasoning
            
            if effective_content:
                if first_token_time is None:
                    first_token_time = time.time() - start_time
                
                total_tokens += estimate_tokens(effective_content)
        
        total_time = time.time() - start_time
        tokens_per_second = total_tokens / total_time if total_time > 0 else 0
        
        # 如果没有收到任何token，记录失败
        # If no tokens received, record failure
        if first_token_time is None:
            first_token_time = -1
            
    except Exception as e:
        print(f"      API错误 / API error: {e}")
        return {
            "first_token_latency": -1,
            "tokens_per_second": -1,
            "total_tokens": 0,
            "total_time": round(time.time() - start_time, 3)
        }
    
    return {
        "first_token_latency": round(first_token_time * 1000, 1) if first_token_time and first_token_time > 0 else -1,
        "tokens_per_second": round(tokens_per_second, 2),
        "total_tokens": int(total_tokens),
        "total_time": round(total_time, 3)
    }


def estimate_tokens(text: str) -> float:
    """估算token数量 / Estimate token count"""
    if not text:
        return 0
    # 简单估算：中文每个字符算0.5 token，英文每个单词约1.3 token
    # Simple estimate: Chinese chars ~0.5 token, English words ~1.3 token
    chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
    other_chars = len(text) - chinese_chars
    return chinese_chars * 0.5 + other_chars / 4


def test_anthropic(client, model: str, prompt: str) -> dict:
    """测试 Anthropic 格式模型 / Test Anthropic format models"""
    start_time = time.time()
    first_token_time = None
    total_tokens = 0
    
    try:
        with client.messages.stream(
            model=model,
            max_tokens=1024,
            messages=[{"role": "user", "content": prompt}]
        ) as stream:
            for chunk in stream:
                if chunk.type == "content_block_delta":
                    if first_token_time is None:
                        first_token_time = time.time() - start_time
                    
                    if hasattr(chunk, 'delta') and hasattr(chunk.delta, 'text'):
                        total_tokens += estimate_tokens(chunk.delta.text)
        
        total_time = time.time() - start_time
        tokens_per_second = total_tokens / total_time if total_time > 0 else 0
        
        if first_token_time is None:
            first_token_time = -1
            
    except Exception as e:
        print(f"      API错误 / API error: {e}")
        return {
            "first_token_latency": -1,
            "tokens_per_second": -1,
            "total_tokens": 0,
            "total_time": round(time.time() - start_time, 3)
        }
    
    return {
        "first_token_latency": round(first_token_time * 1000, 1) if first_token_time and first_token_time > 0 else -1,
        "tokens_per_second": round(tokens_per_second, 2),
        "total_tokens": int(total_tokens),
        "total_time": round(total_time, 3)
    }


def run_tests(config_path: str = "config.yaml"):
    """运行所有测试 / Run all tests"""
    # 读取配置 / Load config
    with open(config_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    
    test_prompt = config.get("test_prompt", "请介绍一下你自己，要求不少于200字")
    models_config = config.get("models", [])
    
    results = []
    
    for model_config in models_config:
        base_url = model_config.get("base_url", "")
        model_type = model_config.get("type", "")
        models = model_config.get("models", [])
        api_key = model_config.get("api_key", "")
        
        print(f"\n{'='*50}")
        print(f"测试 / Testing: {base_url} ({model_type})")
        print(f"{'='*50}")
        
        if model_type == "openai-chat":
            if OpenAI is None:
                print("错误: 请安装 openai 库 / Error: install openai (pip install openai)")
                continue
            
            client = OpenAI(base_url=base_url, api_key=api_key)
            
            for model_name in models:
                print(f"\n测试模型 / Testing model: {model_name}")
                try:
                    result = test_openai(client, model_name, test_prompt)
                    results.append({
                        "base_url": base_url,
                        "type": model_type,
                        "model": model_name,
                        "first_token_latency": result["first_token_latency"],
                        "tokens_per_second": result["tokens_per_second"]
                    })
                    print(f"  首次响应 / First token: {result['first_token_latency']}ms")
                    print(f"  Token/s: {result['tokens_per_second']}")
                    print(f"  总Token / Total: {result['total_tokens']}, 耗时 / Time: {result['total_time']}s")
                except Exception as e:
                    print(f"  错误 / Error: {e}")
                    results.append({
                        "base_url": base_url,
                        "type": model_type,
                        "model": model_name,
                        "first_token_latency": -1,
                        "tokens_per_second": -1
                    })
        
        elif model_type == "anthropic":
            if anthropic is None:
                print("错误: 请安装 anthropic 库 / Error: install anthropic (pip install anthropic)")
                continue
            
            client = anthropic.Anthropic(base_url=base_url, api_key=api_key)
            
            for model_name in models:
                print(f"\n测试模型 / Testing model: {model_name}")
                try:
                    result = test_anthropic(client, model_name, test_prompt)
                    results.append({
                        "base_url": base_url,
                        "type": model_type,
                        "model": model_name,
                        "first_token_latency": result["first_token_latency"],
                        "tokens_per_second": result["tokens_per_second"]
                    })
                    print(f"  首次响应 / First token: {result['first_token_latency']}ms")
                    print(f"  Token/s: {result['tokens_per_second']}")
                    print(f"  总Token / Total: {result['total_tokens']}, 耗时 / Time: {result['total_time']}s")
                except Exception as e:
                    print(f"  错误 / Error: {e}")
                    results.append({
                        "base_url": base_url,
                        "type": model_type,
                        "model": model_name,
                        "first_token_latency": -1,
                        "tokens_per_second": -1
                    })
        
        else:
            print(f"未知类型 / Unknown type: {model_type}")
    
    # 保存CSV / Save CSV
    output_file = "results.csv"
    with open(output_file, 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=["base_url", "type", "model", "first_token_latency", "tokens_per_second"])
        writer.writeheader()
        writer.writerows(results)
    
    print(f"\n{'='*50}")
    print(f"结果已保存到 / Results saved to: {output_file}")
    print(f"{'='*50}")
    
    # 打印汇总表 / Print summary table
    print("\n汇总结果 / Summary:")
    print("-" * 80)
    print(f"{'base_url':<35} {'type':<15} {'model':<25} {'首响(ms)/FT':<14} {'Token/s':<10}")
    print("-" * 80)
    for r in results:
        print(f"{r['base_url']:<35} {r['type']:<15} {r['model']:<25} {r['first_token_latency']:<14} {r['tokens_per_second']:<10}")


if __name__ == "__main__":
    config_file = sys.argv[1] if len(sys.argv) > 1 else "config.yaml"
    run_tests(config_file)