72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
"""
|
||
llm_client.py — Тонкий HTTP-клиент к LLM API.
|
||
|
||
Только отправить промпт → получить ответ. Никакой бизнес-логики.
|
||
Бизнес-логика (промпты, парсинг ответа) → extractor.py.
|
||
|
||
Использует httpx с HTTP/2 — api.aillm.ru за ddos-guard, требует HTTP/2.
|
||
h2 (pure Python) в requirements.txt отдельной строкой.
|
||
|
||
Переменные окружения:
|
||
LLM_API_URL — URL API (по умолчанию https://api.aillm.ru/v1/chat/completions)
|
||
LLM_API_KEY — ключ авторизации
|
||
"""
|
||
|
||
import os
|
||
import httpx
|
||
|
||
# ── Конфигурация ────────────────────────────────────────────────
|
||
|
||
DEFAULT_URL = "https://api.aillm.ru/v1/chat/completions"
|
||
DEFAULT_MODEL = "gpt-oss-120b"
|
||
|
||
|
||
def ask(prompt: str, model: str = None, max_tokens: int = 8000) -> dict:
|
||
"""
|
||
Отправить промпт к LLM API.
|
||
|
||
Returns:
|
||
{"text": "...", "model": "...", "usage": {...}} или {"error": "..."}
|
||
"""
|
||
api_url = os.getenv("LLM_API_URL", DEFAULT_URL)
|
||
api_key = os.getenv("LLM_API_KEY", "")
|
||
|
||
if not api_key:
|
||
return {"error": "LLM_API_KEY not set"}
|
||
|
||
model = model or DEFAULT_MODEL
|
||
|
||
payload = {
|
||
"model": model,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"max_tokens": max_tokens,
|
||
"temperature": 0.1,
|
||
}
|
||
|
||
headers = {
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
try:
|
||
with httpx.Client(http2=True, timeout=120) as client:
|
||
resp = client.post(api_url, json=payload, headers=headers)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
choice = data.get("choices", [{}])[0]
|
||
message = choice.get("message", {})
|
||
return {
|
||
"text": message.get("content", ""),
|
||
"model": data.get("model", model),
|
||
"usage": data.get("usage", {}),
|
||
}
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
http_ver = e.response.http_version if hasattr(e.response, 'http_version') else '?'
|
||
return {"error": f"LLM HTTP {e.response.status_code} (HTTP/{http_ver}): {e.response.text[:500]}"}
|
||
except httpx.TimeoutException:
|
||
return {"error": "LLM timeout"}
|
||
except Exception as e:
|
||
return {"error": f"LLM: {e}"}
|