106 lines
2.9 KiB
Python
106 lines
2.9 KiB
Python
"""Deepseek API 流式快速验证脚本。
|
||
|
||
按照需求：
|
||
- deepseek-chat: max_tokens=8000
|
||
- deepseek-reasoner: max_tokens=65536
|
||
- 开启流式输出，打印每条 data 行便于确认返回。
|
||
|
||
根据用户允许，直接在脚本中硬编码 DeepSeek 测试用 API 信息，避免被仓库中其他
|
||
默认配置（如 kimi）覆盖。请勿在生产使用。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from typing import Optional
|
||
|
||
import httpx
|
||
|
||
|
||
DEEPSEEK_BASE_URL = "https://api.deepseek.com"
|
||
DEEPSEEK_API_KEY = "sk-3457fbc33f0b4aefb2ce1d3101bb2341"
|
||
DEEPSEEK_CHAT_MODEL = "deepseek-chat"
|
||
DEEPSEEK_REASONER_MODEL = "deepseek-reasoner"
|
||
|
||
|
||
def build_headers(api_key: str) -> dict[str, str]:
|
||
return {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
|
||
async def stream_call(
|
||
*,
|
||
name: str,
|
||
base_url: str,
|
||
api_key: str,
|
||
model: str,
|
||
max_tokens: int,
|
||
prompt: Optional[str] = None,
|
||
) -> None:
|
||
"""向指定模型发起流式 chat/completions 请求并打印 data 行。"""
|
||
url = base_url.rstrip("/") + "/chat/completions"
|
||
payload = {
|
||
"model": model,
|
||
"stream": True,
|
||
"max_tokens": max_tokens,
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": prompt
|
||
or f"这是 {name} 模型的流式测试，请用一句话自我介绍。",
|
||
}
|
||
],
|
||
}
|
||
|
||
print(f"\n=== {name} ===")
|
||
print(f"POST {url}")
|
||
async with httpx.AsyncClient(http2=True, timeout=120) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
url,
|
||
json=payload,
|
||
headers=build_headers(api_key),
|
||
) as resp:
|
||
print(f"status: {resp.status_code}")
|
||
if resp.status_code != 200:
|
||
body = await resp.aread()
|
||
print("error body:", body.decode(errors="ignore"))
|
||
return
|
||
async for line in resp.aiter_lines():
|
||
if not line:
|
||
continue
|
||
if line.startswith("data:"):
|
||
data = line[5:].strip()
|
||
if data == "[DONE]":
|
||
print("[DONE]")
|
||
break
|
||
print(data)
|
||
else:
|
||
# 兼容潜在的非 data 行（例如心跳）
|
||
print(line)
|
||
|
||
|
||
async def main() -> None:
|
||
await stream_call(
|
||
name="deepseek-chat (max_tokens=8000)",
|
||
base_url=DEEPSEEK_BASE_URL,
|
||
api_key=DEEPSEEK_API_KEY,
|
||
model=DEEPSEEK_CHAT_MODEL,
|
||
max_tokens=8000,
|
||
)
|
||
|
||
await stream_call(
|
||
name="deepseek-reasoner (max_tokens=65536)",
|
||
base_url=DEEPSEEK_BASE_URL,
|
||
api_key=DEEPSEEK_API_KEY,
|
||
model=DEEPSEEK_REASONER_MODEL,
|
||
max_tokens=65536,
|
||
prompt="你是一个思考模型，请简述测试状态并结束。",
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|