107 lines
4.3 KiB
Python
107 lines
4.3 KiB
Python
"""
|
|
LLM-клиент для DrHider.
|
|
|
|
Используется в TwoPassObfuscator._scan_llm_ner() для обнаружения
|
|
имён, компаний, адресов и паспортных данных через LLM API.
|
|
|
|
Интерфейс:
|
|
client = LLMClient()
|
|
result = client.complete(prompt) # str
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
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)
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Инициализировать клиент. httpx импортируется лениво.
|
|
|
|
Ведёт кумулятивные счётчики за всё время жизни клиента:
|
|
tokens_prompt — суммарные входные токены
|
|
tokens_completion — суммарные выходные токены
|
|
llm_sec — суммарное время LLM-запросов (сек)
|
|
|
|
Плюс флаг текущей активности (для live-таймера в UI):
|
|
llm_active — True пока идёт LLM-вызов
|
|
llm_started — time.time() начала текущего LLM-вызова
|
|
"""
|
|
self._httpx = httpx
|
|
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 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._httpx.post(
|
|
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,
|
|
)
|
|
r.raise_for_status()
|
|
|
|
# Время запроса (реальное время HTTP-вызова LLM)
|
|
self.llm_sec += r.elapsed.total_seconds()
|
|
|
|
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
|