mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
test: 添加 OpenAI chat completions 测试脚本
This commit is contained in:
140
tests/test_chat.py
Normal file
140
tests/test_chat.py
Normal file
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Chat Completions Test Script
|
||||
|
||||
模拟 AsyncOpenAI Python SDK 客户端发送请求,包含完整的请求头和请求体。
|
||||
|
||||
Usage:
|
||||
python test_chat.py
|
||||
python test_chat.py --stream
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
|
||||
import requests
|
||||
|
||||
BASE_URL = "http://localhost:8084/v1"
|
||||
API_KEY = "sk-PCr5oXZNKb9HcyzYqTIMvr8zXsIBK3WS"
|
||||
|
||||
|
||||
def build_headers(api_key: str) -> dict[str, str]:
|
||||
"""模拟 AsyncOpenAI/Python 2.14.0 客户端请求头"""
|
||||
return {
|
||||
"accept": "application/json",
|
||||
"content-type": "application/json",
|
||||
"user-agent": "AsyncOpenAI/Python 2.14.0",
|
||||
"x-stainless-lang": "python",
|
||||
"x-stainless-package-version": "2.14.0",
|
||||
"x-stainless-os": "Windows",
|
||||
"x-stainless-arch": "other:amd64",
|
||||
"x-stainless-runtime": "CPython",
|
||||
"x-stainless-runtime-version": "3.12.11",
|
||||
"x-stainless-async": "async:asyncio",
|
||||
"x-stainless-retry-count": "0",
|
||||
"x-stainless-read-timeout": "120",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
|
||||
def build_payload(stream: bool = False) -> dict:
|
||||
"""构造请求体"""
|
||||
return {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "hi"
|
||||
}
|
||||
],
|
||||
"model": "gpt-5.2",
|
||||
"max_tokens": 2048,
|
||||
"stream": False,
|
||||
"temperature": 0.5
|
||||
}
|
||||
|
||||
|
||||
def send_sync(api_key: str) -> dict:
|
||||
"""发送非流式请求"""
|
||||
url = f"{BASE_URL}/chat/completions"
|
||||
headers = build_headers(api_key)
|
||||
payload = build_payload(stream=False)
|
||||
|
||||
print(f"[sync] POST {url}")
|
||||
print(f"[sync] model={payload['model']}")
|
||||
print("-" * 60)
|
||||
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=120)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
|
||||
usage = data.get("usage", {})
|
||||
|
||||
print(content)
|
||||
print("-" * 60)
|
||||
print(f"Model: {data.get('model')}")
|
||||
print(
|
||||
f"Usage: prompt={usage.get('prompt_tokens', 0)}, "
|
||||
f"completion={usage.get('completion_tokens', 0)}, "
|
||||
f"total={usage.get('total_tokens', 0)}"
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
def send_stream(api_key: str) -> None:
|
||||
"""发送流式请求"""
|
||||
url = f"{BASE_URL}/chat/completions"
|
||||
headers = build_headers(api_key)
|
||||
payload = build_payload(stream=True)
|
||||
|
||||
print(f"[stream] POST {url}")
|
||||
print(f"[stream] model={payload['model']}")
|
||||
print("-" * 60)
|
||||
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=120, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
full_content = ""
|
||||
for line in response.iter_lines(decode_unicode=True):
|
||||
if not line or not line.startswith("data: "):
|
||||
continue
|
||||
data = line[len("data: "):]
|
||||
if data.strip() == "[DONE]":
|
||||
break
|
||||
try:
|
||||
chunk = json.loads(data)
|
||||
delta = chunk.get("choices", [{}])[0].get("delta", {})
|
||||
content = delta.get("content", "")
|
||||
if content:
|
||||
print(content, end="", flush=True)
|
||||
full_content += content
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
print()
|
||||
print("-" * 60)
|
||||
print(f"Total length: {len(full_content)} chars")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
stream = "--stream" in sys.argv
|
||||
|
||||
try:
|
||||
if stream:
|
||||
send_stream(API_KEY)
|
||||
else:
|
||||
send_sync(API_KEY)
|
||||
except requests.exceptions.HTTPError as e:
|
||||
print(f"HTTP Error: {e}", file=sys.stderr)
|
||||
if e.response is not None:
|
||||
print(f"Response: {e.response.text}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user