253 lines
7.7 KiB
Python
Executable File
253 lines
7.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
OpenRouter API 客户端
|
|
===================
|
|
|
|
支持 GPT-4o, Claude, Qwen 等多个模型
|
|
|
|
作者: xialab
|
|
日期: 2026-04-27
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import List, Dict, Tuple
|
|
import time
|
|
|
|
# 加载环境变量
|
|
from dotenv import load_dotenv
|
|
PROJECT_ROOT = Path("/home/xzq/gene-family-mas-v2")
|
|
load_dotenv(PROJECT_ROOT / ".env")
|
|
|
|
# 尝试导入 openai
|
|
try:
|
|
from openai import OpenAI
|
|
except ImportError:
|
|
print("[错误] 请先安装 openai: pip install openai")
|
|
sys.exit(1)
|
|
|
|
|
|
class OpenRouterClient:
|
|
"""OpenRouter API 客户端"""
|
|
|
|
# 可用模型列表
|
|
MODELS = {
|
|
"gpt-4o": "openai/gpt-4o",
|
|
"gpt-4o-mini": "openai/gpt-4o-mini",
|
|
"claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",
|
|
"claude-3-opus": "anthropic/claude-3-opus",
|
|
"qwen-2.5-72b": "qwen/qwen-2.5-72b-instruct",
|
|
"gemini-1.5-pro": "google/gemini-pro-1.5",
|
|
"llama-3.1-70b": "meta-llama/llama-3.1-70b-instruct",
|
|
}
|
|
|
|
def __init__(self):
|
|
self.api_key = os.getenv("OPENROUTER_API_KEY")
|
|
self.base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
|
|
self.referer = os.getenv("OPENROUTER_REFERER", "https://localhost")
|
|
self.title = os.getenv("OPENROUTER_TITLE", "GeneFamily-RAG")
|
|
|
|
if not self.api_key:
|
|
raise ValueError("OPENROUTER_API_KEY not found in environment")
|
|
|
|
self.client = OpenAI(
|
|
base_url=self.base_url,
|
|
api_key=self.api_key,
|
|
default_headers={
|
|
"HTTP-Referer": self.referer,
|
|
"X-Title": self.title,
|
|
}
|
|
)
|
|
print(f"[✓] OpenRouter client initialized")
|
|
print(f" Base URL: {self.base_url}")
|
|
|
|
def query(self, question: str, model: str = "gpt-4o", temperature: float = 0.3) -> Tuple[str, float, Dict]:
|
|
"""
|
|
调用OpenRouter API
|
|
|
|
Args:
|
|
question: 问题文本
|
|
model: 模型名称 (gpt-4o, claude-3.5-sonnet, 等)
|
|
temperature: 温度参数
|
|
|
|
Returns:
|
|
(answer, latency, metadata)
|
|
"""
|
|
start_time = time.time()
|
|
|
|
# 获取完整模型ID
|
|
model_id = self.MODELS.get(model, model)
|
|
|
|
try:
|
|
response = self.client.chat.completions.create(
|
|
model=model_id,
|
|
messages=[
|
|
{
|
|
"role": "system",
|
|
"content": "You are a gene family analysis expert. Answer accurately based on your knowledge."
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": question
|
|
}
|
|
],
|
|
temperature=temperature,
|
|
max_tokens=2048,
|
|
)
|
|
|
|
latency = time.time() - start_time
|
|
answer = response.choices[0].message.content
|
|
|
|
metadata = {
|
|
"model": model,
|
|
"model_id": model_id,
|
|
"prompt_tokens": response.usage.prompt_tokens,
|
|
"completion_tokens": response.usage.completion_tokens,
|
|
"total_tokens": response.usage.total_tokens,
|
|
}
|
|
|
|
return answer, latency, metadata
|
|
|
|
except Exception as e:
|
|
latency = time.time() - start_time
|
|
return f"[错误] {str(e)}", latency, {"error": str(e)}
|
|
|
|
def query_multiple(self, question: str, models: List[str] = None) -> Dict[str, Dict]:
|
|
"""
|
|
使用多个模型查询同一问题
|
|
|
|
Args:
|
|
question: 问题文本
|
|
models: 模型列表，默认 [gpt-4o, claude-3.5-sonnet, qwen-2.5-72b]
|
|
|
|
Returns:
|
|
各模型的回答结果
|
|
"""
|
|
if models is None:
|
|
models = ["gpt-4o", "claude-3.5-sonnet", "qwen-2.5-72b"]
|
|
|
|
results = {}
|
|
print(f"\n问题: {question}\n")
|
|
print("=" * 70)
|
|
|
|
for model in models:
|
|
print(f"\n[{model}] 查询中...")
|
|
answer, latency, metadata = self.query(question, model)
|
|
|
|
results[model] = {
|
|
"answer": answer,
|
|
"latency": latency,
|
|
"metadata": metadata
|
|
}
|
|
|
|
print(f" 耗时: {latency:.2f}s")
|
|
print(f" Token: {metadata.get('total_tokens', 'N/A')}")
|
|
print(f" 回答: {answer[:100]}...")
|
|
|
|
return results
|
|
|
|
|
|
def test_single():
|
|
"""测试单个问题"""
|
|
client = OpenRouterClient()
|
|
|
|
question = "拟南芥中SPL基因家族有哪些成员？"
|
|
|
|
# 使用 gpt-4o-mini (更便宜)
|
|
answer, latency, metadata = client.query(question, model="gpt-4o-mini")
|
|
|
|
print(f"\n问题: {question}")
|
|
print(f"模型: {metadata['model_id']}")
|
|
print(f"耗时: {latency:.2f}s")
|
|
print(f"Token: {metadata['total_tokens']}")
|
|
print(f"\n回答:\n{answer}")
|
|
|
|
|
|
def test_multiple():
|
|
"""测试多个模型对比"""
|
|
client = OpenRouterClient()
|
|
|
|
question = "SPL基因家族在植物发育中起什么作用？"
|
|
models = ["gpt-4o", "claude-3.5-sonnet"]
|
|
|
|
results = client.query_multiple(question, models)
|
|
|
|
# 保存结果
|
|
import json
|
|
output_file = PROJECT_ROOT / "evaluation" / "results" / "openrouter_comparison.json"
|
|
with open(output_file, 'w', encoding='utf-8') as f:
|
|
json.dump(results, f, ensure_ascii=False, indent=2)
|
|
|
|
print(f"\n结果已保存: {output_file}")
|
|
|
|
|
|
def run_benchmark():
|
|
"""运行完整基准测试"""
|
|
client = OpenRouterClient()
|
|
|
|
# 测试问题
|
|
questions = [
|
|
"拟南芥中SPL基因家族有哪些成员？",
|
|
"SPL转录因子如何调控植物开花时间？",
|
|
"SPL基因家族在单子叶和双子叶植物中的扩张模式",
|
|
]
|
|
|
|
models = ["gpt-4o", "claude-3.5-sonnet", "qwen-2.5-72b"]
|
|
|
|
all_results = []
|
|
|
|
for q in questions:
|
|
print(f"\n{'='*70}")
|
|
print(f"问题: {q}")
|
|
print('='*70)
|
|
|
|
for model in models:
|
|
answer, latency, metadata = client.query(q, model)
|
|
|
|
result = {
|
|
"question": q,
|
|
"model": model,
|
|
"answer": answer,
|
|
"latency": latency,
|
|
"tokens": metadata.get('total_tokens', 0)
|
|
}
|
|
all_results.append(result)
|
|
|
|
print(f"\n[{model}] {latency:.2f}s | {metadata.get('total_tokens', 0)} tokens")
|
|
|
|
# 保存结果
|
|
import json
|
|
from datetime import datetime
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
output_file = PROJECT_ROOT / "evaluation" / "results" / f"openrouter_benchmark_{timestamp}.json"
|
|
|
|
with open(output_file, 'w', encoding='utf-8') as f:
|
|
json.dump(all_results, f, ensure_ascii=False, indent=2)
|
|
|
|
print(f"\n{'='*70}")
|
|
print(f"基准测试完成! 结果保存: {output_file}")
|
|
|
|
# 统计
|
|
avg_latency = sum(r['latency'] for r in all_results) / len(all_results)
|
|
total_tokens = sum(r['tokens'] for r in all_results)
|
|
print(f"平均响应时间: {avg_latency:.2f}s")
|
|
print(f"总Token消耗: {total_tokens}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description='OpenRouter API 测试')
|
|
parser.add_argument('--mode', choices=['single', 'multiple', 'benchmark'],
|
|
default='single', help='测试模式')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.mode == 'single':
|
|
test_single()
|
|
elif args.mode == 'multiple':
|
|
test_multiple()
|
|
elif args.mode == 'benchmark':
|
|
run_benchmark()
|