Files
drhider/drhider/llm_client.py
T

65 lines
2.3 KiB
Python

"""
LLM-клиент для DrHider.
Используется в TwoPassObfuscator._scan_llm_ner() для обнаружения
имён, компаний, адресов и паспортных данных через LLM API.
Интерфейс:
client = LLMClient()
result = client.complete(prompt) # str
"""
import os
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 импортируется лениво."""
self._httpx = httpx
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")
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()
return r.json()["choices"][0]["message"]["content"]