292 lines
10 KiB
Python
292 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Тестирование LLM-моделей для AI-ассистента «Зам» (АС «Платформа ОПОРА РОССИИ»).
|
||
|
||
Запуск (Windows PowerShell):
|
||
$env:OPENROUTER_API_KEY = "sk-or-..."
|
||
py -3 test_models.py
|
||
|
||
Свой список моделей:
|
||
py -3 test_models.py --models "openai/gpt-4o-mini,anthropic/claude-3.5-haiku"
|
||
|
||
Результат: таблица в консоли + отчёт report_<дата>.md
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
from datetime import datetime
|
||
|
||
import requests
|
||
|
||
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
|
||
MODELS_URL = "https://openrouter.ai/api/v1/models"
|
||
|
||
DEFAULT_MODELS = [
|
||
"openai/gpt-4o-mini",
|
||
"anthropic/claude-3.5-haiku",
|
||
"google/gemini-2.0-flash-001",
|
||
"qwen/qwen-2.5-72b-instruct",
|
||
"deepseek/deepseek-chat",
|
||
"mistralai/mistral-small-24b-instruct-2501",
|
||
]
|
||
|
||
SYSTEM = (
|
||
"Ты — ассистент «Зам» платформы ОПОРА РОССИИ. "
|
||
"Отвечай по-русски, кратко и по делу. Термины: «точка 0», «текущий срез», "
|
||
"«прирост», «динамика», «красная зона», «вклад округа»."
|
||
)
|
||
|
||
TOOLS = [
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "create_task",
|
||
"description": "Создать задачу в трекере",
|
||
"parameters": {
|
||
"type": "object",
|
||
"properties": {
|
||
"title": {"type": "string", "description": "Название задачи"},
|
||
"due_date": {"type": "string", "description": "Срок в формате YYYY-MM-DD"},
|
||
"priority": {"type": "string", "enum": ["низкий", "средний", "высокий"]},
|
||
"assignee": {"type": "string", "description": "Ответственный (регион или роль)"},
|
||
},
|
||
"required": ["title", "due_date", "priority", "assignee"],
|
||
},
|
||
},
|
||
}
|
||
]
|
||
|
||
TESTS = [
|
||
{
|
||
"name": "ru_generation",
|
||
"desc": "Русский текст: еженедельный фокус целей",
|
||
"messages": [
|
||
{"role": "system", "content": SYSTEM},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"Составь сообщение «Еженедельный фокус целей» для Донецкой Народной "
|
||
"Республики. Прирост от точки 0: +136. Округ ЮФО: +365, с динамикой "
|
||
"12/12 регионов. Добавь рекомендацию «что делать»."
|
||
),
|
||
},
|
||
],
|
||
},
|
||
{
|
||
"name": "function_calling",
|
||
"desc": "Вызов инструмента create_task",
|
||
"messages": [
|
||
{"role": "system", "content": SYSTEM},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"Создай задачу: «Подготовить отчёт по приросту», срок 2026-10-01, "
|
||
"приоритет высокий, ответственный — регион ДНР."
|
||
),
|
||
},
|
||
],
|
||
"tools": TOOLS,
|
||
},
|
||
{
|
||
"name": "json_output",
|
||
"desc": "Структурированный JSON",
|
||
"messages": [
|
||
{"role": "system", "content": SYSTEM + " Отвечай строго в формате JSON, без пояснений."},
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"Верни JSON с полями region, point0, current, growth для региона ДНР: "
|
||
"точка 0 = 1000, текущий срез = 1136."
|
||
),
|
||
},
|
||
],
|
||
},
|
||
]
|
||
|
||
|
||
def load_pricing():
|
||
"""Возвращает {model_id: (prompt_price, completion_price)} в $ за 1 токен."""
|
||
try:
|
||
r = requests.get(MODELS_URL, timeout=30)
|
||
r.raise_for_status()
|
||
data = r.json().get("data", [])
|
||
pricing = {}
|
||
for m in data:
|
||
p = m.get("pricing", {})
|
||
try:
|
||
pricing[m["id"]] = (float(p.get("prompt", 0)), float(p.get("completion", 0)))
|
||
except (TypeError, ValueError):
|
||
pass
|
||
return pricing
|
||
except Exception as e:
|
||
print(f"[warn] не удалось загрузить прайс: {e}")
|
||
return {}
|
||
|
||
|
||
def call_model(api_key, model, test):
|
||
payload = {
|
||
"model": model,
|
||
"messages": test["messages"],
|
||
"temperature": 0.3,
|
||
"max_tokens": 700,
|
||
}
|
||
if "tools" in test:
|
||
payload["tools"] = test["tools"]
|
||
payload["tool_choice"] = "auto"
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
"HTTP-Referer": "https://opora.my-dpr.ru",
|
||
"X-Title": "OPORA Zam LLM test",
|
||
}
|
||
|
||
start = time.perf_counter()
|
||
try:
|
||
r = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=120)
|
||
latency = time.perf_counter() - start
|
||
if r.status_code != 200:
|
||
return {"ok": False, "error": f"HTTP {r.status_code}: {r.text[:300]}", "latency": latency}
|
||
data = r.json()
|
||
except Exception as e:
|
||
return {"ok": False, "error": str(e), "latency": time.perf_counter() - start}
|
||
|
||
choice = (data.get("choices") or [{}])[0]
|
||
msg = choice.get("message", {})
|
||
usage = data.get("usage", {}) or {}
|
||
return {
|
||
"ok": True,
|
||
"latency": latency,
|
||
"content": msg.get("content") or "",
|
||
"tool_calls": msg.get("tool_calls") or [],
|
||
"prompt_tokens": usage.get("prompt_tokens", 0),
|
||
"completion_tokens": usage.get("completion_tokens", 0),
|
||
"model_returned": data.get("model", model),
|
||
}
|
||
|
||
|
||
def check_function_call(result):
|
||
"""Проверяет корректность вызова create_task."""
|
||
if not result.get("ok"):
|
||
return "ошибка"
|
||
calls = result.get("tool_calls") or []
|
||
if not calls:
|
||
return "нет вызова"
|
||
try:
|
||
fn = calls[0]["function"]
|
||
if fn["name"] != "create_task":
|
||
return f"не та функция: {fn['name']}"
|
||
args = json.loads(fn["arguments"])
|
||
need = {"title", "due_date", "priority", "assignee"}
|
||
missing = need - set(args)
|
||
if missing:
|
||
return f"нет полей: {', '.join(sorted(missing))}"
|
||
if args.get("due_date") != "2026-10-01":
|
||
return f"срок: {args.get('due_date')}"
|
||
return "OK"
|
||
except Exception as e:
|
||
return f"ошибка разбора: {e}"
|
||
|
||
|
||
def check_json(result):
|
||
if not result.get("ok"):
|
||
return "ошибка"
|
||
txt = (result.get("content") or "").strip()
|
||
txt = txt.replace("```json", "").replace("```", "").strip()
|
||
try:
|
||
obj = json.loads(txt)
|
||
need = {"region", "point0", "current", "growth"}
|
||
missing = need - set(obj)
|
||
if missing:
|
||
return f"нет полей: {', '.join(sorted(missing))}"
|
||
if int(obj.get("growth", 0)) != 136:
|
||
return f"growth={obj.get('growth')}"
|
||
return "OK"
|
||
except Exception as e:
|
||
return f"не JSON: {e}"
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--models", default=",".join(DEFAULT_MODELS))
|
||
ap.add_argument("--out", default=None)
|
||
args = ap.parse_args()
|
||
|
||
api_key = os.environ.get("OPENROUTER_API_KEY")
|
||
if not api_key:
|
||
print("Ошибка: задайте переменную окружения OPENROUTER_API_KEY")
|
||
sys.exit(1)
|
||
|
||
models = [m.strip() for m in args.models.split(",") if m.strip()]
|
||
pricing = load_pricing()
|
||
|
||
rows = []
|
||
for model in models:
|
||
print(f"\n=== {model} ===")
|
||
row = {"model": model}
|
||
for test in TESTS:
|
||
res = call_model(api_key, model, test)
|
||
if not res["ok"]:
|
||
print(f" [{test['name']}] ОШИБКА: {res['error']}")
|
||
row[test["name"]] = {"status": "ошибка", "latency": res["latency"]}
|
||
continue
|
||
|
||
if test["name"] == "function_calling":
|
||
status = check_function_call(res)
|
||
elif test["name"] == "json_output":
|
||
status = check_json(res)
|
||
else:
|
||
status = "OK" if len(res["content"]) > 40 else "короткий ответ"
|
||
|
||
pt, ct = res["prompt_tokens"], res["completion_tokens"]
|
||
price = pricing.get(model, (0, 0))
|
||
cost = pt * price[0] + ct * price[1]
|
||
|
||
row[test["name"]] = {
|
||
"status": status,
|
||
"latency": res["latency"],
|
||
"prompt_tokens": pt,
|
||
"completion_tokens": ct,
|
||
"cost_usd": cost,
|
||
"sample": res["content"][:200],
|
||
}
|
||
print(f" [{test['name']}] {status} | {res['latency']:.2f}s | "
|
||
f"{pt}+{ct} ток | ${cost:.6f}")
|
||
|
||
rows.append(row)
|
||
|
||
# Отчёт
|
||
ts = datetime.now().strftime("%Y-%m-%d_%H%M")
|
||
out = args.out or f"report_{ts}.md"
|
||
lines = [f"# Тест LLM-моделей — {datetime.now():%Y-%m-%d %H:%M}", ""]
|
||
lines.append("| Модель | ru_generation | function_calling | json_output | Ср. латентность | Стоимость теста |")
|
||
lines.append("|---|---|---|---|---|---|")
|
||
for row in rows:
|
||
lat = [row[t["name"]]["latency"] for t in TESTS if row.get(t["name"], {}).get("latency")]
|
||
avg_lat = sum(lat) / len(lat) if lat else 0
|
||
cost = sum(row[t["name"]].get("cost_usd", 0) for t in TESTS if row.get(t["name"]))
|
||
cells = []
|
||
for t in TESTS:
|
||
d = row.get(t["name"], {})
|
||
cells.append(d.get("status", "—"))
|
||
lines.append(f"| {row['model']} | " + " | ".join(cells) +
|
||
f" | {avg_lat:.2f}s | ${cost:.6f} |")
|
||
lines.append("")
|
||
lines.append("## Примеры ответов (ru_generation)")
|
||
for row in rows:
|
||
d = row.get("ru_generation", {})
|
||
if d.get("sample"):
|
||
lines.append(f"\n### {row['model']}\n\n{d['sample']}")
|
||
|
||
with open(out, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(lines))
|
||
print(f"\nОтчёт сохранён: {out}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|