82 lines
3.4 KiB
Python
82 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
from urllib.parse import urljoin
|
|
|
|
import httpx
|
|
|
|
|
|
SYSTEM_PROMPT = """Ты интерактивный Telegram-помощник для поездок и пересечения границы Беларусь-Польша.
|
|
Отвечай кратко, по-русски, практично. Используй текущие данные очередей из контекста.
|
|
Не выдавай юридические гарантии: если вопрос о правилах, визах, детях, документах или запретах,
|
|
советуй проверить официальный источник и указывай, что правила могут измениться.
|
|
""".strip()
|
|
|
|
|
|
class AiAssistant:
|
|
def __init__(self, base_url: str | None, api_key: str | None, model: str, timeout: float) -> None:
|
|
self.base_url = base_url.rstrip("/") if base_url else None
|
|
self.api_key = api_key
|
|
self.model = model
|
|
self.timeout = timeout
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return bool(self.base_url and self.api_key)
|
|
|
|
async def ask(self, question: str, border_context: str) -> str:
|
|
if not self.enabled:
|
|
return "AI сейчас не настроен. Добавь AI_BASE_URL и AI_API_KEY в переменные окружения."
|
|
|
|
payload = {
|
|
"model": self.model,
|
|
"messages": [
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": f"Текущие данные:\n{border_context}\n\nВопрос пользователя:\n{question}"},
|
|
],
|
|
"temperature": 0.2,
|
|
}
|
|
headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
|
|
|
|
async with httpx.AsyncClient(timeout=self.timeout, follow_redirects=True) as client:
|
|
response = await self._post_chat_completion(client, payload, headers)
|
|
response.raise_for_status()
|
|
|
|
data = response.json()
|
|
return _extract_answer(data)
|
|
|
|
async def _post_chat_completion(
|
|
self,
|
|
client: httpx.AsyncClient,
|
|
payload: dict,
|
|
headers: dict[str, str],
|
|
) -> httpx.Response:
|
|
assert self.base_url is not None
|
|
candidates = [self.base_url]
|
|
if not self.base_url.endswith("/v1/chat/completions"):
|
|
candidates.append(urljoin(f"{self.base_url}/", "v1/chat/completions"))
|
|
|
|
last_response: httpx.Response | None = None
|
|
for url in candidates:
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
if response.status_code not in {404, 405}:
|
|
return response
|
|
last_response = response
|
|
assert last_response is not None
|
|
return last_response
|
|
|
|
|
|
def _extract_answer(data: dict) -> str:
|
|
choices = data.get("choices") or []
|
|
if choices:
|
|
message = choices[0].get("message") or {}
|
|
content = message.get("content")
|
|
if isinstance(content, str) and content.strip():
|
|
return content.strip()
|
|
|
|
for key in ("answer", "response", "content", "text"):
|
|
value = data.get(key)
|
|
if isinstance(value, str) and value.strip():
|
|
return value.strip()
|
|
|
|
return "AI вернул ответ в неожиданном формате. Попробуй переформулировать вопрос."
|