140 lines
6.1 KiB
Python
140 lines
6.1 KiB
Python
"""
|
|
LLM-клиент для Obfuscator.
|
|
|
|
Используется в TwoPassObfuscator._scan_llm_ner() для обнаружения
|
|
имён, компаний, адресов и паспортных данных через LLM API.
|
|
|
|
Интерфейс:
|
|
client = LLMClient()
|
|
result = client.complete(prompt) # str
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
|
|
|
|
class LLMClient:
|
|
"""Клиент для вызова LLM API (aillm.ru / OpenAI-совместимый).
|
|
|
|
Использует переменные окружения:
|
|
LLM_API_KEY или LLM_KEY — ключ API
|
|
LLM_URL — URL эндпоинта (по умолчанию https://api.aillm.ru/v1/chat/completions)
|
|
LLM_MODEL — модель (по умолчанию gpt-oss-120b)
|
|
LLM_RETRIES — число повторов после первой попытки (по умолчанию 3)
|
|
LLM_RETRY_DELAY — базовая задержка backoff в секундах (по умолчанию 2)
|
|
"""
|
|
|
|
def __init__(self, retries: Optional[int] = None,
|
|
retry_delay: Optional[float] = None):
|
|
"""Инициализировать клиент. httpx импортируется лениво.
|
|
|
|
Ведёт кумулятивные счётчики за всё время жизни клиента:
|
|
tokens_prompt — суммарные входные токены
|
|
tokens_completion — суммарные выходные токены
|
|
llm_sec — суммарное время LLM-запросов (сек)
|
|
|
|
Плюс флаг текущей активности (для live-таймера в UI):
|
|
llm_active — True пока идёт LLM-вызов
|
|
llm_started — time.time() начала текущего LLM-вызова
|
|
"""
|
|
self._httpx = httpx
|
|
self.retries = int(os.environ.get("LLM_RETRIES", "3")) if retries is None else retries
|
|
self.retry_delay = (
|
|
float(os.environ.get("LLM_RETRY_DELAY", "2.0"))
|
|
if retry_delay is None else retry_delay
|
|
)
|
|
if self.retries < 0 or self.retry_delay < 0:
|
|
raise ValueError("LLM retries and retry delay must be non-negative")
|
|
self.tokens_prompt = 0
|
|
self.tokens_completion = 0
|
|
self.llm_sec = 0.0
|
|
self.llm_active = False
|
|
self.llm_started = 0.0
|
|
|
|
@property
|
|
def tokens_total(self) -> int:
|
|
return self.tokens_prompt + self.tokens_completion
|
|
|
|
def llm_elapsed_now(self) -> float:
|
|
"""Секунд с начала текущего LLM-вызова (0, если LLM не активен)."""
|
|
if self.llm_active:
|
|
return time.time() - self.llm_started
|
|
return 0.0
|
|
|
|
def _post_with_retries(self, url: str, json: dict, headers: dict,
|
|
timeout: int) -> httpx.Response:
|
|
"""Повторить только временные транспортные ошибки и выбранные HTTP-статусы."""
|
|
for attempt in range(self.retries + 1):
|
|
try:
|
|
response = self._httpx.post(
|
|
url, json=json, headers=headers, timeout=timeout
|
|
)
|
|
self.llm_sec += response.elapsed.total_seconds()
|
|
response.raise_for_status()
|
|
return response
|
|
except (self._httpx.TimeoutException, self._httpx.TransportError):
|
|
if attempt >= self.retries:
|
|
raise
|
|
except self._httpx.HTTPStatusError as e:
|
|
# Retry only throttling and server failures; 400/401/403 require
|
|
# a corrected request or credentials and will not improve on retry.
|
|
if e.response.status_code not in (429, 500, 502, 503, 504):
|
|
raise
|
|
if attempt >= self.retries:
|
|
raise
|
|
|
|
# 120 секунд на попытку и backoff означают около 8 минут в худшем случае.
|
|
time.sleep(self.retry_delay * (attempt + 1))
|
|
|
|
def complete(self, prompt: str) -> str:
|
|
"""Отправить промпт в LLM и вернуть текст ответа.
|
|
|
|
Args:
|
|
prompt: Текст промпта (инструкция + данные для анализа)
|
|
|
|
Returns:
|
|
Текстовый ответ модели (обычно JSON-строка)
|
|
|
|
Raises:
|
|
httpx.HTTPError: при ошибке HTTP
|
|
KeyError: при неожиданном формате ответа
|
|
"""
|
|
# Ключ API: сначала LLM_KEY, затем LLM_API_KEY (для совместимости)
|
|
key = os.environ.get("LLM_KEY") or os.environ.get("LLM_API_KEY", "")
|
|
|
|
# URL и модель с значениями по умолчанию
|
|
url = os.environ.get("LLM_URL", "https://api.aillm.ru/v1/chat/completions")
|
|
model = os.environ.get("LLM_MODEL", "gpt-oss-120b")
|
|
|
|
# Пометим LLM-вызов как активный (для live-таймера в UI)
|
|
self.llm_active = True
|
|
self.llm_started = time.time()
|
|
try:
|
|
r = self._post_with_retries(
|
|
url,
|
|
json={
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 8000,
|
|
"temperature": 0.1,
|
|
},
|
|
headers={
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
timeout=120,
|
|
)
|
|
|
|
data = r.json()
|
|
# Счётчики токенов из usage (могут отсутствовать у некоторых прокси)
|
|
usage = data.get("usage", {}) or {}
|
|
self.tokens_prompt += int(usage.get("prompt_tokens", 0) or 0)
|
|
self.tokens_completion += int(usage.get("completion_tokens", 0) or 0)
|
|
|
|
return data["choices"][0]["message"]["content"]
|
|
finally:
|
|
self.llm_active = False
|