Files
contracts-flask/site/services/llm_client.py
T
2026-07-15 11:23:13 +04:00

75 lines
2.8 KiB
Python

"""LLM Client — протокол + httpx-реализация + fake для тестов.
План decoupling Ф2:
- LLMClient.complete(prompt) -> str — единственный метод
- Парсинг JSON остаётся у вызывающего (call_llm / _call_llm_classify)
- HttpxLLMClient — продакшен
- FakeLLMClient — тесты (возвращает сохранённые ответы)
"""
from typing import Protocol
class LLMClient(Protocol):
"""Протокол LLM-клиента. Один метод — сырой вызов."""
def complete(self, prompt: str) -> str:
"""Отправить prompt в LLM, вернуть raw_text."""
...
class HttpxLLMClient:
"""Продакшен-клиент через httpx → api.aillm.ru."""
def __init__(self, url: str, key: str, model: str = "gpt-oss-120b",
max_tokens: int = 8000, temperature: float = 0.1, timeout: int = 120):
self.url = url
self.key = key
self.model = model
self.max_tokens = max_tokens
self.temperature = temperature
self.timeout = timeout
def complete(self, prompt: str) -> str:
import httpx
payload = {
"model": self.model,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": self.max_tokens,
"temperature": self.temperature,
}
with httpx.Client(http2=True, timeout=self.timeout, verify=True) as client:
resp = client.post(
self.url,
json=payload,
headers={
"Authorization": f"Bearer {self.key}",
"Content-Type": "application/json",
},
)
resp.raise_for_status()
data = resp.json()
return data["choices"][0]["message"]["content"]
class FakeLLMClient:
"""Тестовый клиент — возвращает сохранённые ответы из словаря."""
def __init__(self, responses: dict = None):
self.responses = responses or {}
self.calls: list[str] = [] # история вызовов для проверок
def add(self, prompt_key: str, response: str):
"""Зарегистрировать ответ для конкретного prompt_key."""
self.responses[prompt_key] = response
def complete(self, prompt: str) -> str:
self.calls.append(prompt)
# Ищем точное совпадение
if prompt in self.responses:
return self.responses[prompt]
# Ищем по частичному ключу
for key, resp in self.responses.items():
if key in prompt:
return resp
# Fallback
return self.responses.get("default", '{"error": "no response registered"}')