66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
|
|
def post_chat(base_url: str, model: str, prompt: str) -> dict:
|
|
url = f"{base_url.rstrip('/')}/v1/chat/completions"
|
|
payload = {
|
|
"model": model,
|
|
"messages": [
|
|
{"role": "system", "content": "Отвечай кратко и по делу."},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"temperature": 0.2,
|
|
"max_tokens": 256,
|
|
}
|
|
request = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(request, timeout=120) as response:
|
|
return json.loads(response.read().decode("utf-8"))
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Smoke test an OpenAI-compatible chat endpoint.")
|
|
parser.add_argument("--base-url", default="http://docker-gpu.cin.su:8000")
|
|
parser.add_argument("--model", default="qwen3-4b-instruct")
|
|
parser.add_argument(
|
|
"--prompt",
|
|
default="Объясни в одном предложении, зачем нужен реестр локальных моделей.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
response = post_chat(args.base_url, args.model, args.prompt)
|
|
except urllib.error.URLError as exc:
|
|
print(f"Request failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
choices = response.get("choices") or []
|
|
if not choices:
|
|
print("Response has no choices.", file=sys.stderr)
|
|
print(json.dumps(response, ensure_ascii=False, indent=2))
|
|
return 1
|
|
|
|
message = choices[0].get("message") or {}
|
|
content = (message.get("content") or "").strip()
|
|
if not content:
|
|
print("Response content is empty.", file=sys.stderr)
|
|
print(json.dumps(response, ensure_ascii=False, indent=2))
|
|
return 1
|
|
|
|
print(content)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|