Files
xiaozhi-esp32-server/main/xiaozhi-server/performance_tester/performance_tester_llm.py
T

545 lines
21 KiB
Python

import asyncio
import logging
import os
import statistics
import time
import concurrent.futures
from typing import Dict, Optional
import yaml
import aiohttp
from tabulate import tabulate
from core.utils.llm import create_instance as create_llm_instance
from config.settings import load_config
# 设置全局日志级别为 WARNING,抑制 INFO 级别日志
logging.basicConfig(level=logging.WARNING)
description = "大语言模型性能测试"
class LLMPerformanceTester:
def __init__(self):
self.config = load_config()
# 使用更符合智能体场景的测试内容,包含系统提示词
self.system_prompt = self._load_system_prompt()
self.test_sentences = self.config.get("module_test", {}).get(
"test_sentences",
[
"你好,我今天心情不太好,能安慰一下我吗?",
"帮我查一下明天的天气如何?",
"我想听一个有趣的故事,你能给我讲一个吗?",
"现在几点了?今天是星期几?",
"我想设置一个明天早上8点的闹钟提醒我开会",
],
)
self.results = {}
def _load_system_prompt(self) -> str:
"""加载系统提示词"""
try:
prompt_file = os.path.join(
os.path.dirname(os.path.dirname(__file__)), "agent-base-prompt.txt"
)
with open(prompt_file, "r", encoding="utf-8") as f:
content = f.read()
# 替换模板变量为测试值
content = content.replace(
"{{base_prompt}}", "你是小智,一个聪明可爱的AI助手"
)
content = content.replace(
"{{emojiList}}", "😀,😃,😄,😁,😊,😍,🤔,😮,😱,😢,😭,😴,😵,🤗,🙄"
)
content = content.replace("{{current_time}}", "2024年8月17日 12:30:45")
content = content.replace("{{today_date}}", "2024年8月17日")
content = content.replace("{{today_weekday}}", "星期六")
content = content.replace("{{lunar_date}}", "甲辰年七月十四")
content = content.replace("{{local_address}}", "北京市")
content = content.replace("{{weather_info}}", "今天晴,25-32℃")
return content
except Exception as e:
print(f"无法加载系统提示词文件: {e}")
return "你是小智,一个聪明可爱的AI助手。请用温暖友善的语气回复用户。"
def _collect_response_sync(self, llm, messages, llm_name, sentence_start):
"""同步收集响应数据的辅助方法"""
chunks = []
first_token_received = False
first_token_time = None
try:
response_generator = llm.response("perf_test", messages)
chunk_count = 0
for chunk in response_generator:
chunk_count += 1
# 每处理一定数量的chunk就检查一下是否应该中断
if chunk_count % 10 == 0:
# 通过检查当前线程是否被标记为中断来提前退出
import threading
if (
threading.current_thread().ident
!= threading.main_thread().ident
):
# 如果不是主线程,检查是否应该停止
pass
# 检查chunk是否包含错误信息
chunk_str = str(chunk)
if (
"异常" in chunk_str
or "错误" in chunk_str
or "502" in chunk_str.lower()
):
error_msg = chunk_str.lower()
print(f"{llm_name} 响应包含错误信息: {error_msg}")
# 抛出一个包含错误信息的异常
raise Exception(chunk_str)
if not first_token_received and chunk.strip() != "":
first_token_time = time.time() - sentence_start
first_token_received = True
print(f"{llm_name} 首个 Token: {first_token_time:.3f}s")
chunks.append(chunk)
except Exception as e:
# 更详细的错误信息
error_msg = str(e).lower()
print(f"{llm_name} 响应收集异常: {error_msg}")
# 对于502错误或网络错误,直接抛出异常让上层处理
if (
"502" in error_msg
or "bad gateway" in error_msg
or "error code: 502" in error_msg
or "异常" in str(e)
or "错误" in str(e)
):
raise e
# 对于其他错误,可以返回部分结果
return chunks, first_token_time
return chunks, first_token_time
async def _check_ollama_service(self, base_url: str, model_name: str) -> bool:
"""异步检查 Ollama 服务状态"""
async with aiohttp.ClientSession() as session:
try:
async with session.get(f"{base_url}/api/version") as response:
if response.status != 200:
print(f"Ollama 服务未启动或无法访问: {base_url}")
return False
async with session.get(f"{base_url}/api/tags") as response:
if response.status == 200:
data = await response.json()
models = data.get("models", [])
if not any(model["name"] == model_name for model in models):
print(
f"Ollama 模型 {model_name} 未找到,请先使用 `ollama pull {model_name}` 下载"
)
return False
else:
print("无法获取 Ollama 模型列表")
return False
return True
except Exception as e:
print(f"无法连接到 Ollama 服务: {str(e)}")
return False
async def _test_single_sentence(
self, llm_name: str, llm, sentence: str
) -> Optional[Dict]:
"""测试单个句子的性能"""
try:
print(f"{llm_name} 开始测试: {sentence[:20]}...")
sentence_start = time.time()
first_token_received = False
first_token_time = None
# 构建包含系统提示词的消息
messages = [
{"role": "system", "content": self.system_prompt},
{"role": "user", "content": sentence},
]
# 使用asyncio.wait_for进行超时控制
try:
loop = asyncio.get_event_loop()
with concurrent.futures.ThreadPoolExecutor() as executor:
# 创建响应收集任务
future = executor.submit(
self._collect_response_sync,
llm,
messages,
llm_name,
sentence_start,
)
# 使用asyncio.wait_for实现超时控制
try:
response_chunks, first_token_time = await asyncio.wait_for(
asyncio.wrap_future(future), timeout=10.0
)
except asyncio.TimeoutError:
print(f"{llm_name} 测试超时(10秒),跳过")
# 强制取消future
future.cancel()
# 等待一小段时间确保线程池任务能够响应取消
try:
await asyncio.wait_for(
asyncio.wrap_future(future), timeout=1.0
)
except (
asyncio.TimeoutError,
concurrent.futures.CancelledError,
Exception,
):
# 忽略所有异常,确保程序继续执行
pass
return None
except Exception as timeout_error:
print(f"{llm_name} 处理异常: {timeout_error}")
return None
response_time = time.time() - sentence_start
print(f"{llm_name} 完成响应: {response_time:.3f}s")
return {
"name": llm_name,
"type": "llm",
"first_token_time": first_token_time,
"response_time": response_time,
}
except Exception as e:
error_msg = str(e).lower()
# 检查是否为502错误或网络错误
if (
"502" in error_msg
or "bad gateway" in error_msg
or "error code: 502" in error_msg
):
print(f"{llm_name} 遇到502错误,跳过测试")
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": "502网络错误",
}
print(f"{llm_name} 句子测试失败: {str(e)}")
return None
async def _test_llm(self, llm_name: str, config: Dict) -> Dict:
"""异步测试单个 LLM 性能"""
try:
# 对于 Ollama,跳过 api_key 检查并进行特殊处理
if llm_name == "Ollama":
base_url = config.get("base_url", "http://localhost:11434")
model_name = config.get("model_name")
if not model_name:
print("Ollama 未配置 model_name")
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": "网络错误",
}
if not await self._check_ollama_service(base_url, model_name):
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": "网络错误",
}
else:
if "api_key" in config and any(
x in config["api_key"] for x in ["你的", "placeholder", "sk-xxx"]
):
print(f"跳过未配置的 LLM: {llm_name}")
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": "配置错误",
}
# 获取实际类型(兼容旧配置)
module_type = config.get("type", llm_name)
llm = create_llm_instance(module_type, config)
# 统一使用 UTF-8 编码
test_sentences = [
s.encode("utf-8").decode("utf-8") for s in self.test_sentences
]
# 创建所有句子的测试任务
sentence_tasks = []
for sentence in test_sentences:
sentence_tasks.append(
self._test_single_sentence(llm_name, llm, sentence)
)
# 并发执行所有句子测试,并处理可能的异常
sentence_results = await asyncio.gather(
*sentence_tasks, return_exceptions=True
)
# 处理结果,过滤掉异常和None值
valid_results = []
for result in sentence_results:
if isinstance(result, dict) and result is not None:
valid_results.append(result)
elif isinstance(result, Exception):
error_msg = str(result).lower()
if "502" in error_msg or "bad gateway" in error_msg:
print(f"{llm_name} 遇到502错误,跳过该句子测试")
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": "502网络错误",
}
else:
print(f"{llm_name} 句子测试异常: {result}")
if not valid_results:
print(f"{llm_name} 无有效数据,可能遇到网络问题或配置错误")
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": "网络错误",
}
# 检查有效结果数量,如果太少则认为测试失败
if len(valid_results) < len(test_sentences) * 0.3: # 至少要有30%的成功率
print(
f"{llm_name} 成功测试句子过少({len(valid_results)}/{len(test_sentences)}),可能网络不稳定或接口有问题"
)
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": "网络错误",
}
first_token_times = [
r["first_token_time"]
for r in valid_results
if r.get("first_token_time")
]
response_times = [r["response_time"] for r in valid_results]
# 过滤异常数据(超出3个标准差的数据)
if len(response_times) > 1:
mean = statistics.mean(response_times)
stdev = statistics.stdev(response_times)
filtered_times = [t for t in response_times if t <= mean + 3 * stdev]
else:
filtered_times = response_times
return {
"name": llm_name,
"type": "llm",
"avg_response": sum(response_times) / len(response_times),
"avg_first_token": (
sum(first_token_times) / len(first_token_times)
if first_token_times
else 0
),
"success_rate": f"{len(valid_results)}/{len(test_sentences)}",
"errors": 0,
}
except Exception as e:
error_msg = str(e).lower()
if "502" in error_msg or "bad gateway" in error_msg:
print(f"LLM {llm_name} 遇到502错误,跳过测试")
else:
print(f"LLM {llm_name} 测试失败: {str(e)}")
error_type = "网络错误"
if "timeout" in str(e).lower():
error_type = "超时连接"
return {
"name": llm_name,
"type": "llm",
"errors": 1,
"error_type": error_type,
}
def _print_results(self):
"""打印测试结果"""
print("\n" + "=" * 50)
print("LLM 性能测试结果")
print("=" * 50)
if not self.results:
print("没有可用的测试结果")
return
headers = ["模型名称", "平均响应时间(s)", "首Token时间(s)", "成功率", "状态"]
table_data = []
# 收集所有数据并分类
valid_results = []
error_results = []
for name, data in self.results.items():
if data["errors"] == 0:
# 正常结果
avg_response = f"{data['avg_response']:.3f}"
avg_first_token = (
f"{data['avg_first_token']:.3f}"
if data["avg_first_token"] > 0
else "-"
)
success_rate = data.get("success_rate", "N/A")
status = "✅ 正常"
# 保存用于排序的值
first_token_value = (
data["avg_first_token"]
if data["avg_first_token"] > 0
else float("inf")
)
valid_results.append(
{
"name": name,
"avg_response": avg_response,
"avg_first_token": avg_first_token,
"success_rate": success_rate,
"status": status,
"sort_key": first_token_value,
}
)
else:
# 错误结果
avg_response = "-"
avg_first_token = "-"
success_rate = "0/5"
# 获取具体错误类型
error_type = data.get("error_type", "网络错误")
status = f"❌ {error_type}"
error_results.append(
[name, avg_response, avg_first_token, success_rate, status]
)
# 按首Token时间升序排序
valid_results.sort(key=lambda x: x["sort_key"])
# 将排序后的有效结果转换为表格数据
for result in valid_results:
table_data.append(
[
result["name"],
result["avg_response"],
result["avg_first_token"],
result["success_rate"],
result["status"],
]
)
# 将错误结果添加到表格数据末尾
table_data.extend(error_results)
print(tabulate(table_data, headers=headers, tablefmt="grid"))
print("\n测试说明:")
print("- 测试内容:包含完整系统提示词的智能体对话场景")
print("- 超时控制:单个请求最大等待时间为10秒")
print("- 错误处理:自动跳过502错误和网络异常的模型")
print("- 成功率:成功响应的句子数量/总测试句子数量")
print("\n测试完成!")
async def run(self):
"""执行全量异步测试"""
print("开始筛选可用 LLM 模块...")
# 创建所有测试任务
all_tasks = []
# LLM 测试任务
if self.config.get("LLM") is not None:
for llm_name, config in self.config.get("LLM", {}).items():
# 检查配置有效性
if llm_name == "CozeLLM":
if any(x in config.get("bot_id", "") for x in ["你的"]) or any(
x in config.get("user_id", "") for x in ["你的"]
):
print(f"LLM {llm_name} 未配置 bot_id/user_id,已跳过")
continue
elif "api_key" in config and any(
x in config["api_key"] for x in ["你的", "placeholder", "sk-xxx"]
):
print(f"LLM {llm_name} 未配置 api_key,已跳过")
continue
# 对于 Ollama,先检查服务状态
if llm_name == "Ollama":
base_url = config.get("base_url", "http://localhost:11434")
model_name = config.get("model_name")
if not model_name:
print("Ollama 未配置 model_name")
continue
if not await self._check_ollama_service(base_url, model_name):
continue
print(f"添加 LLM 测试任务: {llm_name}")
all_tasks.append(self._test_llm(llm_name, config))
print(f"\n找到 {len(all_tasks)} 个可用 LLM 模块")
print("\n开始并发测试所有模块...\n")
# 并发执行所有测试任务,但为每个任务设置独立超时
async def test_with_timeout(task, timeout=30):
"""为每个测试任务添加超时保护"""
try:
return await asyncio.wait_for(task, timeout=timeout)
except asyncio.TimeoutError:
print(f"测试任务超时({timeout}秒),跳过")
return {
"name": "Unknown",
"type": "llm",
"errors": 1,
"error_type": "超时连接",
}
except Exception as e:
print(f"测试任务异常: {str(e)}")
return {
"name": "Unknown",
"type": "llm",
"errors": 1,
"error_type": "网络错误",
}
# 为每个任务包装超时保护
protected_tasks = [test_with_timeout(task) for task in all_tasks]
# 并发执行所有测试任务
all_results = await asyncio.gather(*protected_tasks, return_exceptions=True)
# 处理结果
for result in all_results:
if isinstance(result, dict):
if result.get("errors") == 0:
self.results[result["name"]] = result
else:
# 即使有错误也记录,用于显示失败状态
if result.get("name") != "Unknown":
self.results[result["name"]] = result
elif isinstance(result, Exception):
print(f"测试结果处理异常: {str(result)}")
# 打印结果
print("\n生成测试报告...")
self._print_results()
async def main():
tester = LLMPerformanceTester()
await tester.run()
if __name__ == "__main__":
asyncio.run(main())