105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
豆包（火山引擎 Ark）API 快速测试
|
|||
|
|
用法: python3 test-doubao-api.py
|
|||
|
|
"""
|
|||
|
|
import json, time, sys
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import requests
|
|||
|
|
except ImportError:
|
|||
|
|
print("需要安装 requests: pip install requests")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
API_KEY = "18480c26-ebcd-4263-8a6f-48359b8bd65d"
|
|||
|
|
BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
|
|||
|
|
MODEL = "doubao-seed-2-0-pro-260215"
|
|||
|
|
|
|||
|
|
headers = {
|
|||
|
|
"Content-Type": "application/json",
|
|||
|
|
"Authorization": f"Bearer {API_KEY}"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
print("=" * 50)
|
|||
|
|
print(" 豆包 API 测试")
|
|||
|
|
print("=" * 50)
|
|||
|
|
|
|||
|
|
# ---- 测试 1: 非流式 ----
|
|||
|
|
print("\n【测试 1】非流式调用")
|
|||
|
|
payload = {
|
|||
|
|
"model": MODEL,
|
|||
|
|
"messages": [{"role": "user", "content": "你好，请用一句话介绍你自己"}],
|
|||
|
|
"stream": False
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
start = time.time()
|
|||
|
|
try:
|
|||
|
|
resp = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload, timeout=30)
|
|||
|
|
elapsed = (time.time() - start) * 1000
|
|||
|
|
print(f"HTTP 状态码: {resp.status_code}")
|
|||
|
|
print(f"耗时: {elapsed:.0f} ms")
|
|||
|
|
if resp.status_code == 200:
|
|||
|
|
data = resp.json()
|
|||
|
|
choices = data.get("choices", [])
|
|||
|
|
if choices:
|
|||
|
|
print(f"回复: {choices[0]['message']['content']}")
|
|||
|
|
usage = data.get("usage", {})
|
|||
|
|
print(f"Token: prompt={usage.get('prompt_tokens',0)}, completion={usage.get('completion_tokens',0)}, total={usage.get('total_tokens',0)}")
|
|||
|
|
print(f"模型: {data.get('model', '?')}")
|
|||
|
|
else:
|
|||
|
|
print(f"错误: {resp.text[:500]}")
|
|||
|
|
except Exception as e:
|
|||
|
|
elapsed = (time.time() - start) * 1000
|
|||
|
|
print(f"请求失败 ({elapsed:.0f}ms): {e}")
|
|||
|
|
|
|||
|
|
# ---- 测试 2: 流式 ----
|
|||
|
|
print("\n【测试 2】流式调用 (SSE)")
|
|||
|
|
payload_stream = {
|
|||
|
|
"model": MODEL,
|
|||
|
|
"messages": [{"role": "user", "content": "简单说一下健康饮食的三个要点，每个要点一句话"}],
|
|||
|
|
"stream": True
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
start2 = time.time()
|
|||
|
|
first_token_time = None
|
|||
|
|
full_content = ""
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
resp2 = requests.post(f"{BASE_URL}/chat/completions", headers=headers, json=payload_stream, stream=True, timeout=30)
|
|||
|
|
print(f"HTTP 状态码: {resp2.status_code}")
|
|||
|
|
if resp2.status_code == 200:
|
|||
|
|
print("流式输出: ", end="", flush=True)
|
|||
|
|
for line in resp2.iter_lines(decode_unicode=True):
|
|||
|
|
if not line:
|
|||
|
|
continue
|
|||
|
|
if line.startswith("data: "):
|
|||
|
|
json_str = line[6:].strip()
|
|||
|
|
if json_str == "[DONE]":
|
|||
|
|
break
|
|||
|
|
try:
|
|||
|
|
evt = json.loads(json_str)
|
|||
|
|
delta = evt.get("choices", [{}])[0].get("delta", {}).get("content", "")
|
|||
|
|
if delta:
|
|||
|
|
if first_token_time is None:
|
|||
|
|
first_token_time = time.time()
|
|||
|
|
full_content += delta
|
|||
|
|
print(delta, end="", flush=True)
|
|||
|
|
except json.JSONDecodeError:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
total_time = (time.time() - start2) * 1000
|
|||
|
|
ttft = (first_token_time - start2) * 1000 if first_token_time else 0
|
|||
|
|
print(f"\n---")
|
|||
|
|
print(f"首字时间 (TTFT): {ttft:.0f} ms")
|
|||
|
|
print(f"总耗时: {total_time:.0f} ms")
|
|||
|
|
print(f"回复长度: {len(full_content)} 字符")
|
|||
|
|
else:
|
|||
|
|
print(f"错误: {resp2.text[:500]}")
|
|||
|
|
except Exception as e:
|
|||
|
|
elapsed2 = (time.time() - start2) * 1000
|
|||
|
|
print(f"\n请求失败 ({elapsed2:.0f}ms): {e}")
|
|||
|
|
|
|||
|
|
print("\n" + "=" * 50)
|
|||
|
|
print(" 测试完成")
|
|||
|
|
print("=" * 50)
|