86 lines
2.5 KiB
Go
86 lines
2.5 KiB
Go
package main
|
|
|
|
import (
|
|
"flag"
|
|
"fmt"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/acmestudio/llm-api-benchmark-tool/config"
|
|
"github.com/acmestudio/llm-api-benchmark-tool/benchmark"
|
|
"github.com/acmestudio/llm-api-benchmark-tool/report"
|
|
"github.com/acmestudio/llm-api-benchmark-tool/logger"
|
|
)
|
|
|
|
func main() {
|
|
// 解析命令行参数
|
|
configPath := flag.String("config", "config.yaml", "配置文件路径")
|
|
outputDir := flag.String("output", "reports", "报告输出目录")
|
|
debugMode := flag.Bool("debug", false, "启用调试模式")
|
|
noColor := flag.Bool("no-color", false, "禁用彩色日志输出")
|
|
flag.Parse()
|
|
|
|
// 配置日志级别
|
|
if *debugMode {
|
|
logger.SetLevel(logger.DEBUG)
|
|
} else {
|
|
logger.SetLevel(logger.INFO)
|
|
}
|
|
|
|
// 配置日志颜色
|
|
logger.SetUseColors(!*noColor)
|
|
|
|
logger.Info("LLM API基准测试工具启动")
|
|
logger.Info("版本: v0.1.0")
|
|
logger.Debug("调试模式已启用")
|
|
|
|
// 加载配置
|
|
logger.Info("正在加载配置文件: %s", *configPath)
|
|
cfg, err := config.LoadConfig(*configPath)
|
|
if err != nil {
|
|
logger.Fatal("加载配置失败: %v", err)
|
|
}
|
|
logger.Info("配置加载成功")
|
|
logger.Debug("API端点: %s", cfg.API.Endpoint)
|
|
logger.Debug("模型: %s", cfg.API.Model)
|
|
logger.Debug("并发步骤: %v", cfg.Concurrency.Steps)
|
|
logger.Debug("每步持续时间: %d秒", cfg.Concurrency.DurationPerStep)
|
|
|
|
// 创建输出目录
|
|
if err := os.MkdirAll(*outputDir, 0755); err != nil {
|
|
logger.Fatal("创建输出目录失败: %v", err)
|
|
}
|
|
logger.Debug("输出目录已创建: %s", *outputDir)
|
|
|
|
// 初始化基准测试
|
|
logger.Info("正在初始化基准测试...")
|
|
bench, err := benchmark.NewBenchmark(cfg)
|
|
if err != nil {
|
|
logger.Fatal("初始化基准测试失败: %v", err)
|
|
}
|
|
logger.Info("基准测试初始化成功")
|
|
|
|
// 运行基准测试
|
|
logger.Info("开始运行LLM API基准测试...")
|
|
startTime := time.Now()
|
|
results, err := bench.Run()
|
|
if err != nil {
|
|
logger.Fatal("运行基准测试失败: %v", err)
|
|
}
|
|
duration := time.Since(startTime)
|
|
logger.Info("基准测试完成，总耗时: %v", duration)
|
|
|
|
// 生成报告
|
|
logger.Info("正在生成性能报告...")
|
|
reportTime := time.Now().Format("2006-01-02_15-04-05")
|
|
reportPath := fmt.Sprintf("%s/report_%s", *outputDir, reportTime)
|
|
|
|
if err := report.GenerateReport(results, cfg, reportPath); err != nil {
|
|
logger.Fatal("生成报告失败: %v", err)
|
|
}
|
|
|
|
logger.Info("基准测试完成！报告已生成: %s", reportPath)
|
|
logger.Info("HTML报告: %s/report.html", reportPath)
|
|
logger.Info("CSV数据: %s/results.csv", reportPath)
|
|
}
|