443 lines
14 KiB
Python
443 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
多 AI 模型调用脚本 (Universal Architect Engine v4.0)
|
||
|
||
支持双中转站 API:
|
||
- 中转站 GPT: gpt-5.2-codex 等
|
||
- 中转站 Gemini: gemini-2.5-pro, gemini-2.5-flash 等
|
||
|
||
环境变量:
|
||
AI_API_KEY - 中转站 API 密钥 (GPT 模型)
|
||
AI_BASE_URL - 中转站地址 (默认: https://api.gbro.site)
|
||
GEMINI_API_KEY - 中转站 API 密钥 (Gemini 模型，可与 AI_API_KEY 相同)
|
||
GEMINI_BASE_URL - Gemini 中转站地址 (默认: 与 AI_BASE_URL 相同)
|
||
|
||
使用方法:
|
||
python multi_ai.py --model gpt-5.2-codex --task audit --file code.py
|
||
python multi_ai.py --model gemini-2.5-pro --task frontend --prompt "生成登录组件"
|
||
python multi_ai.py --list-models
|
||
python multi_ai.py --test
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import argparse
|
||
import requests
|
||
import json
|
||
|
||
# ============================================
|
||
# 配置
|
||
# ============================================
|
||
|
||
# 默认中转站地址
|
||
DEFAULT_BASE_URL = "https://api.gbro.site"
|
||
|
||
# 模型到 API 源的映射
|
||
MODEL_API_SOURCE = {
|
||
# GPT 模型 → GPT 中转站
|
||
"gpt-5.2-codex": "gpt",
|
||
"gpt-5.2": "gpt",
|
||
"gpt-5.1": "gpt",
|
||
"gpt-5.1-codex": "gpt",
|
||
"gpt-5.1-codex-mini": "gpt",
|
||
"gpt-5.1-codex-max": "gpt",
|
||
|
||
# Gemini 模型 → Gemini 中转站
|
||
"gemini-2.5-pro": "gemini",
|
||
"gemini-2.5-flash": "gemini",
|
||
"gemini-3-pro-preview": "gemini",
|
||
"gemini-3-pro-preview-thinking": "gemini",
|
||
"gemini-2.5-flash-image": "gemini",
|
||
"gemini-2.5-flash-image-preview": "gemini",
|
||
|
||
# Claude 模型 → GPT 中转站
|
||
"claude-sonnet-4-5-20250929": "gpt",
|
||
"claude-opus-4-5": "gpt",
|
||
"claude-haiku-4-5-20251001": "gpt",
|
||
}
|
||
|
||
# 模型别名映射
|
||
MODELS = {
|
||
# ===== OpenAI/GPT 系列 =====
|
||
"gpt-5.2-codex": "gpt-5.2-codex",
|
||
"gpt-5.2": "gpt-5.2",
|
||
"gpt-5.1": "gpt-5.1",
|
||
"codex": "gpt-5.2-codex",
|
||
"openai": "gpt-5.2-codex",
|
||
"gpt": "gpt-5.2-codex",
|
||
|
||
# ===== Gemini 系列 =====
|
||
"gemini-2.5-pro": "gemini-2.5-pro",
|
||
"gemini-2.5-flash": "gemini-2.5-flash",
|
||
"gemini-3-pro": "gemini-3-pro-preview",
|
||
"gemini-3-pro-preview": "gemini-3-pro-preview",
|
||
"gemini": "gemini-2.5-flash",
|
||
"flash": "gemini-2.5-flash",
|
||
|
||
# ===== Claude 系列 =====
|
||
"claude-sonnet": "claude-sonnet-4-5-20250929",
|
||
"claude-opus": "claude-opus-4-5",
|
||
"claude-haiku": "claude-haiku-4-5-20251001",
|
||
"claude": "claude-sonnet-4-5-20250929",
|
||
|
||
# ===== 任务别名 =====
|
||
"audit": "gpt-5.2-codex", # 代码审计用 GPT
|
||
"frontend": "gemini-2.5-flash", # 前端生成用 Gemini
|
||
"design": "gemini-2.5-flash", # 设计系统用 Gemini
|
||
"backend": "gpt-5.2-codex", # 后端逻辑用 GPT
|
||
"fast": "gemini-2.5-flash", # 快速任务用 Gemini Flash
|
||
}
|
||
|
||
# 任务提示词模板
|
||
TASK_PROMPTS = {
|
||
"audit": """你是一位首席架构师，请对以下代码进行深度审计。
|
||
|
||
审计维度:
|
||
1. 🔐 Security: 注入预防、鉴权逻辑、敏感数据泄露
|
||
2. ⚡ Performance: N+1 查询、冗余渲染、索引建议
|
||
3. 🔧 Maintainability: 代码耦合度、模块化清晰度
|
||
|
||
请输出结构化报告，区分:
|
||
- 🚫 BLOCKER (必须修改)
|
||
- 💡 ADVICE (改进建议)
|
||
|
||
代码:
|
||
```
|
||
{content}
|
||
```
|
||
|
||
请用中文回复。""",
|
||
|
||
"frontend": """你是一位高级前端架构师，请根据以下需求生成现代化前端代码。
|
||
|
||
要求:
|
||
1. 使用 React + TypeScript + Tailwind CSS
|
||
2. 组件命名具备清晰的业务语义
|
||
3. 包含完整的类型定义
|
||
4. 预留 OpenTelemetry 埋点
|
||
|
||
需求:
|
||
{content}
|
||
|
||
请用中文注释。""",
|
||
|
||
"design-system": """你是一位 UI/UX 设计系统专家，请根据以下项目信息生成 Design Tokens。
|
||
|
||
项目信息:
|
||
{content}
|
||
|
||
请输出 JSON 格式的 Design Tokens，包括:
|
||
1. colors (primary, secondary, semantic)
|
||
2. spacing (unit, scale)
|
||
3. typography (fontFamily, fontSize)
|
||
4. borderRadius
|
||
5. shadows
|
||
|
||
同时提供关键 UI 组件的设计规范描述。""",
|
||
|
||
"backend": """你是一位后端架构专家，请根据以下需求设计后端架构。
|
||
|
||
要求:
|
||
1. 遵循 SOLID 原则和 Clean Architecture
|
||
2. 考虑高并发和幂等性设计
|
||
3. 预留分布式追踪埋点
|
||
4. 包含异常处理框架
|
||
|
||
需求:
|
||
{content}
|
||
|
||
请用中文注释。""",
|
||
|
||
"contract": """你是一位全栈架构师，请根据以下 API 设计生成前后端握手契约。
|
||
|
||
要求:
|
||
1. TypeScript Interfaces 或 Protobuf 格式
|
||
2. 包含 DTO (Data Transfer Object) 定义
|
||
3. 包含错误码枚举
|
||
4. 确保前端组件与后端 100% 类型对齐
|
||
|
||
API 设计:
|
||
{content}
|
||
|
||
请输出完整的 Handshake Contract。""",
|
||
|
||
"general": """{content}"""
|
||
}
|
||
|
||
|
||
def get_config():
|
||
"""获取配置"""
|
||
# GPT 中转站配置
|
||
gpt_key = os.environ.get("AI_API_KEY")
|
||
gpt_url = os.environ.get("AI_BASE_URL") or DEFAULT_BASE_URL
|
||
|
||
# Gemini 中转站配置（可单独设置，也可复用 GPT 配置）
|
||
gemini_key = os.environ.get("GEMINI_API_KEY") or gpt_key
|
||
gemini_url = os.environ.get("GEMINI_BASE_URL") or gpt_url
|
||
|
||
return {
|
||
"gpt_key": gpt_key,
|
||
"gpt_url": gpt_url,
|
||
"gemini_key": gemini_key,
|
||
"gemini_url": gemini_url
|
||
}
|
||
|
||
|
||
def call_model(model: str, prompt: str, temperature: float = 0.3) -> str:
|
||
"""调用模型"""
|
||
config = get_config()
|
||
|
||
# 解析模型名称
|
||
actual_model = MODELS.get(model, model)
|
||
|
||
# 确定使用哪个 API 源
|
||
api_source = MODEL_API_SOURCE.get(actual_model, "gpt")
|
||
|
||
if api_source == "gemini":
|
||
return call_proxy_api(
|
||
actual_model, prompt,
|
||
config["gemini_key"], config["gemini_url"],
|
||
temperature
|
||
)
|
||
else:
|
||
return call_proxy_api(
|
||
actual_model, prompt,
|
||
config["gpt_key"], config["gpt_url"],
|
||
temperature
|
||
)
|
||
|
||
|
||
def call_proxy_api(model: str, prompt: str, api_key: str, base_url: str, temperature: float) -> str:
|
||
"""调用中转站 API (OpenAI 兼容格式)"""
|
||
if not api_key:
|
||
print("❌ 未设置 API Key 环境变量")
|
||
sys.exit(1)
|
||
|
||
url = f"{base_url}/v1/chat/completions"
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {api_key}"
|
||
}
|
||
data = {
|
||
"model": model,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"temperature": temperature
|
||
}
|
||
|
||
response = requests.post(url, headers=headers, json=data, timeout=180)
|
||
response.raise_for_status()
|
||
|
||
result = response.json()
|
||
return result["choices"][0]["message"]["content"]
|
||
|
||
|
||
def list_models():
|
||
"""列出可用模型"""
|
||
config = get_config()
|
||
|
||
print("=" * 60)
|
||
print("📋 可用模型列表")
|
||
print("=" * 60)
|
||
|
||
# 显示 GPT 中转站模型
|
||
print("\n🔶 GPT 中转站模型:")
|
||
print(f" 地址: {config['gpt_url']}")
|
||
if config["gpt_key"]:
|
||
url = f"{config['gpt_url']}/v1/models"
|
||
headers = {"Authorization": f"Bearer {config['gpt_key']}"}
|
||
|
||
try:
|
||
response = requests.get(url, headers=headers, timeout=30)
|
||
response.raise_for_status()
|
||
|
||
models = response.json().get("data", [])
|
||
gpt_models = [m["id"] for m in models if "gpt" in m["id"].lower() or "codex" in m["id"].lower()]
|
||
|
||
for m in sorted(gpt_models)[:10]:
|
||
print(f" - {m}")
|
||
if len(gpt_models) > 10:
|
||
print(f" ... 还有 {len(gpt_models) - 10} 个")
|
||
|
||
except Exception as e:
|
||
print(f" ❌ 获取失败: {e}")
|
||
else:
|
||
print(" ⚠️ 未设置 AI_API_KEY")
|
||
|
||
# 显示 Gemini 中转站模型
|
||
print("\n🔷 Gemini 中转站模型:")
|
||
print(f" 地址: {config['gemini_url']}")
|
||
if config["gemini_key"]:
|
||
url = f"{config['gemini_url']}/v1/models"
|
||
headers = {"Authorization": f"Bearer {config['gemini_key']}"}
|
||
|
||
try:
|
||
response = requests.get(url, headers=headers, timeout=30)
|
||
response.raise_for_status()
|
||
|
||
models = response.json().get("data", [])
|
||
gemini_models = [m["id"] for m in models if "gemini" in m["id"].lower()]
|
||
|
||
for m in sorted(gemini_models)[:10]:
|
||
print(f" - {m}")
|
||
if len(gemini_models) > 10:
|
||
print(f" ... 还有 {len(gemini_models) - 10} 个")
|
||
|
||
except Exception as e:
|
||
print(f" ❌ 获取失败: {e}")
|
||
else:
|
||
print(" ⚠️ 未设置 GEMINI_API_KEY")
|
||
|
||
print("\n" + "=" * 60)
|
||
print("环境变量配置:")
|
||
print(f" AI_API_KEY: {'✅ 已设置' if config['gpt_key'] else '❌ 未设置'} (GPT)")
|
||
print(f" GEMINI_API_KEY: {'✅ 已设置' if config['gemini_key'] else '❌ 未设置'} (Gemini)")
|
||
print("=" * 60)
|
||
|
||
|
||
def test_models():
|
||
"""测试主要模型是否可用"""
|
||
config = get_config()
|
||
|
||
print("=" * 60)
|
||
print("🔍 测试模型可用性")
|
||
print("=" * 60)
|
||
|
||
# 测试 GPT 中转站
|
||
print("\n--- GPT 中转站 ---")
|
||
if config["gpt_key"]:
|
||
for model, name in [("gpt-5.2-codex", "GPT-5.2 Codex")]:
|
||
print(f"\n📝 测试 {name}...")
|
||
try:
|
||
result = call_proxy_api(model, "说你好", config["gpt_key"], config["gpt_url"], 0.1)
|
||
print(f" ✅ 可用: {result[:30]}...")
|
||
except requests.exceptions.HTTPError as e:
|
||
error_msg = str(e)
|
||
if hasattr(e, 'response'):
|
||
error_msg = e.response.text[:100]
|
||
print(f" ❌ 失败: {error_msg}")
|
||
except Exception as e:
|
||
print(f" ❌ 失败: {e}")
|
||
else:
|
||
print(" ⚠️ 未设置 AI_API_KEY")
|
||
|
||
# 测试 Gemini 中转站
|
||
print("\n--- Gemini 中转站 ---")
|
||
if config["gemini_key"]:
|
||
for model, name in [
|
||
("gemini-2.5-flash", "Gemini 2.5 Flash"),
|
||
("gemini-2.5-pro", "Gemini 2.5 Pro"),
|
||
]:
|
||
print(f"\n📝 测试 {name}...")
|
||
try:
|
||
result = call_proxy_api(model, "说你好", config["gemini_key"], config["gemini_url"], 0.1)
|
||
print(f" ✅ 可用: {result[:30]}...")
|
||
except requests.exceptions.HTTPError as e:
|
||
error_msg = str(e)
|
||
if hasattr(e, 'response'):
|
||
try:
|
||
err_json = e.response.json()
|
||
error_msg = err_json.get("error", {}).get("message", str(e))[:80]
|
||
except:
|
||
error_msg = e.response.text[:80]
|
||
print(f" ❌ 失败: {error_msg}")
|
||
except Exception as e:
|
||
print(f" ❌ 失败: {e}")
|
||
else:
|
||
print(" ⚠️ 未设置 GEMINI_API_KEY")
|
||
|
||
print("\n" + "=" * 60)
|
||
print("测试完成")
|
||
print("=" * 60)
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description='多 AI 模型调用工具 (Universal Architect Engine v4.0)',
|
||
formatter_class=argparse.RawDescriptionHelpFormatter
|
||
)
|
||
|
||
parser.add_argument('--model', '-m', default='gpt-5.2-codex',
|
||
help='模型名称 (gpt-5.2-codex, gemini-2.5-pro, etc.)')
|
||
parser.add_argument('--task', '-t', default='general',
|
||
choices=['audit', 'frontend', 'design-system', 'backend', 'contract', 'general'],
|
||
help='任务类型')
|
||
parser.add_argument('--file', '-f', help='输入文件')
|
||
parser.add_argument('--prompt', '-p', help='直接输入提示词')
|
||
parser.add_argument('--output', '-o', help='输出文件')
|
||
parser.add_argument('--temperature', type=float, default=0.3, help='温度参数')
|
||
parser.add_argument('--list-models', action='store_true', help='列出可用模型')
|
||
parser.add_argument('--test', action='store_true', help='测试主要模型可用性')
|
||
parser.add_argument('--json', '-j', action='store_true', help='JSON 格式输出')
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 列出模型
|
||
if args.list_models:
|
||
list_models()
|
||
return
|
||
|
||
# 测试模型
|
||
if args.test:
|
||
test_models()
|
||
return
|
||
|
||
# 获取输入内容
|
||
if args.file:
|
||
with open(args.file, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
elif args.prompt:
|
||
content = args.prompt
|
||
elif not sys.stdin.isatty():
|
||
content = sys.stdin.read()
|
||
else:
|
||
parser.print_help()
|
||
return
|
||
|
||
# 构建提示词
|
||
task_template = TASK_PROMPTS.get(args.task, TASK_PROMPTS["general"])
|
||
full_prompt = task_template.format(content=content)
|
||
|
||
# 输出状态
|
||
actual_model = MODELS.get(args.model, args.model)
|
||
print(f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
|
||
print(f"[Status: Executing {args.task} via {actual_model}]", file=sys.stderr)
|
||
print(f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
|
||
|
||
try:
|
||
result = call_model(args.model, full_prompt, args.temperature)
|
||
|
||
if args.json:
|
||
output = json.dumps({
|
||
"model": actual_model,
|
||
"task": args.task,
|
||
"result": result
|
||
}, ensure_ascii=False, indent=2)
|
||
else:
|
||
output = result
|
||
|
||
# 输出结果
|
||
if args.output:
|
||
with open(args.output, 'w', encoding='utf-8') as f:
|
||
f.write(output)
|
||
print(f"✅ 结果已保存到: {args.output}", file=sys.stderr)
|
||
else:
|
||
print()
|
||
print(output)
|
||
|
||
print(f"\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
|
||
print(f"[Task Complete]", file=sys.stderr)
|
||
print(f"━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", file=sys.stderr)
|
||
|
||
except requests.exceptions.HTTPError as e:
|
||
print(f"❌ API 错误: {e}")
|
||
if hasattr(e, 'response'):
|
||
print(f" 响应: {e.response.text}")
|
||
sys.exit(1)
|
||
except Exception as e:
|
||
print(f"❌ 失败: {e}")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|