101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
"""unifyllm 网关连通性 / tool-calling 冒烟测试。
|
|
|
|
用法(.env 在仓库根, 含 UNIFYLLM_API_KEY):
|
|
.venv/Scripts/python.exe scripts/diag_unifyllm.py [model_id ...]
|
|
|
|
不带参数时测默认 5 个精选模型。走 litellm openai/ 前缀 + api_base,
|
|
与 core/llm.py 线上路径一致。输出一律 ASCII(Windows 控制台 GBK)。
|
|
"""
|
|
import os
|
|
import sys
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
env_file = os.path.join(ROOT, ".env")
|
|
if os.path.exists(env_file):
|
|
for line in open(env_file, encoding="utf-8"):
|
|
line = line.strip()
|
|
if line and not line.startswith("#") and "=" in line:
|
|
k, v = line.split("=", 1)
|
|
os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))
|
|
|
|
os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
|
import litellm # noqa: E402
|
|
|
|
API_BASE = "https://unifyllm.ai/v1"
|
|
DEFAULT_MODELS = [
|
|
"claude-fable-5",
|
|
"claude-opus-4-8",
|
|
"claude-sonnet-4-6",
|
|
"gpt-5.6-sol",
|
|
"gemini-3.1-pro-preview",
|
|
]
|
|
|
|
TOOLS = [{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"description": "Get current weather for a city",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"city": {"type": "string"}},
|
|
"required": ["city"],
|
|
},
|
|
},
|
|
}]
|
|
|
|
|
|
def main() -> int:
|
|
key = os.environ.get("UNIFYLLM_API_KEY")
|
|
if not key:
|
|
print("[FAIL] UNIFYLLM_API_KEY not set (check .env)")
|
|
return 1
|
|
args = sys.argv[1:]
|
|
force_temp = None
|
|
stream = False
|
|
for a in list(args):
|
|
if a.startswith("--temp="):
|
|
force_temp = float(a.split("=", 1)[1])
|
|
args.remove(a)
|
|
elif a == "--stream":
|
|
stream = True
|
|
args.remove(a)
|
|
models = args or DEFAULT_MODELS
|
|
failed = []
|
|
for m in models:
|
|
# litellm 客户端硬拦 gpt-5* 的 temperature != 1
|
|
temp = force_temp if force_temp is not None else (1.0 if m.startswith("gpt-5") else 0.3)
|
|
try:
|
|
kwargs = dict(
|
|
model=f"openai/{m}",
|
|
api_base=API_BASE,
|
|
api_key=key,
|
|
messages=[{"role": "user", "content": "What's the weather in Beijing? Use the tool."}],
|
|
tools=TOOLS,
|
|
temperature=temp,
|
|
timeout=120,
|
|
)
|
|
if stream:
|
|
# 与 core/llm.py chat_stream 同路径:流式 + include_usage,chunk 拼回完整 response
|
|
chunks = list(litellm.completion(**kwargs, stream=True,
|
|
stream_options={"include_usage": True}))
|
|
resp = litellm.stream_chunk_builder(chunks)
|
|
else:
|
|
resp = litellm.completion(**kwargs)
|
|
msg = resp.choices[0].message
|
|
tc = msg.tool_calls or []
|
|
usage = resp.usage
|
|
if tc:
|
|
print(f"[OK] {m}: tool_call={tc[0].function.name}({tc[0].function.arguments.strip()}) "
|
|
f"tokens={usage.prompt_tokens}+{usage.completion_tokens}")
|
|
else:
|
|
text = (msg.content or "")[:80].replace("\n", " ")
|
|
print(f"[WARN] {m}: no tool_call, text={text!r}")
|
|
except Exception as e:
|
|
failed.append(m)
|
|
print(f"[FAIL] {m}: {type(e).__name__}: {str(e)[:200]}")
|
|
return 1 if failed else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|