75 lines
2.8 KiB
Python
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"}')
|