- Add shared proxy module (bot/utils/proxy.py) to eliminate duplicate SOCKS5 parsing - Fix AI client: escape HTML before Markdown→HTML conversion, unify timeouts, make health check optional and disabled by default, handle 429 retries - Fix ai.py: correct forwarded message handling (aiogram 3.x forward_origin), pass relevant summaries via extra_system_content - Fix dialogue.py: only respond to Astra's messages, use system context instead of prompt injection, answer on the limit message before phase transition - Fix voice.py: load Whisper model in thread pool, safe WAV path generation - Improve database.py: composite indexes, Boolean is_active, upsert file_id cache, add context cleanup helper - Update weather.py and yadisk_download.py to use shared proxy connector - Update yadisk.py: validate URL before cache clear, add download size limit, wrap sync file ops in to_thread - Reuse S3 client via lru_cache - Update setup_commands with /aiclear and /aiuser - Update README, Dockerfile (Python 3.11), docker-compose (mount models) - Pin dependency versions, remove unused httpx[socks] - Add basic pytest tests for layout converter and voice normalization
431 lines
14 KiB
Python
431 lines
14 KiB
Python
import asyncio
|
|
import html
|
|
import logging
|
|
import re
|
|
import time
|
|
|
|
import aiohttp
|
|
|
|
from bot.utils.proxy import get_proxy_connector
|
|
from config import (
|
|
OPENROUTER_API_KEY,
|
|
AI_SYSTEM_PROMPT,
|
|
ROUTERAI_API_KEY,
|
|
ROUTERAI_BASE_URL,
|
|
ROUTERAI_MODEL,
|
|
AI_HEALTH_CHECK_ENABLED,
|
|
AI_HEALTH_CHECK_INTERVAL,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
|
|
ROUTERAI_URL = f"{ROUTERAI_BASE_URL}/chat/completions"
|
|
|
|
PAID_NOTICE = "\n\n⚡ Обработано через платный API"
|
|
|
|
_free_models_cache: list[str] = []
|
|
_free_models_cache_time = 0.0
|
|
_free_models_cache_ttl = 3600
|
|
|
|
_working_models_cache: list[str] = []
|
|
_working_models_cache_time = 0.0
|
|
_working_models_cache_ttl = 600
|
|
|
|
|
|
def get_client_timeout(total: int = 60) -> aiohttp.ClientTimeout:
|
|
return aiohttp.ClientTimeout(total=total, sock_connect=15, sock_read=30)
|
|
|
|
|
|
def _md_to_html(text: str) -> str:
|
|
"""Convert a small subset of Markdown to Telegram HTML, escaping raw HTML first."""
|
|
text = html.escape(text)
|
|
|
|
def _pre_repl(match: re.Match) -> str:
|
|
lang = match.group(1)
|
|
code = html.unescape(match.group(2))
|
|
return f'<pre><code class="language-{lang}">{html.escape(code)}</code></pre>' if lang else f"<pre>{html.escape(code)}</pre>"
|
|
|
|
text = re.sub(r"```(\w*)\n(.*?)```", _pre_repl, text, flags=re.DOTALL)
|
|
text = re.sub(r"`([^`]+)`", r"<code>\1</code>", text)
|
|
text = re.sub(r"\*\*(.+?)\*\*", r"<b>\1</b>", text)
|
|
text = re.sub(r"\*(.+?)\*", r"<i>\1</i>", text)
|
|
text = re.sub(r"__(.+?)__", r"<u>\1</u>", text)
|
|
text = re.sub(r"~~(.+?)~~", r"<s>\1</s>", text)
|
|
text = re.sub(r"\[(.+?)\]\((.+?)\)", r'<a href="\2">\1</a>', text)
|
|
return text
|
|
|
|
|
|
async def _fetch_free_models() -> list[str]:
|
|
global _free_models_cache, _free_models_cache_time
|
|
|
|
now = time.time()
|
|
if _free_models_cache and (now - _free_models_cache_time) < _free_models_cache_ttl:
|
|
return _free_models_cache
|
|
|
|
logger.info("Fetching free models from OpenRouter API...")
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {OPENROUTER_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
fallback = [
|
|
"deepseek/deepseek-v4-flash:free",
|
|
"google/gemma-4-26b-a4b-it:free",
|
|
"minimax/minimax-m2.5:free",
|
|
"qwen/qwen3-next-80b-a3b-instruct:free",
|
|
]
|
|
|
|
connector = get_proxy_connector()
|
|
try:
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
async with session.get(
|
|
"https://openrouter.ai/api/v1/models",
|
|
headers=headers,
|
|
timeout=get_client_timeout(30),
|
|
) as response:
|
|
if response.status == 200:
|
|
data = await response.json()
|
|
models = data.get("data", [])
|
|
free_models = [m.get("id", "") for m in models if m.get("id", "").endswith(":free")]
|
|
|
|
if free_models:
|
|
_free_models_cache = free_models
|
|
_free_models_cache_time = now
|
|
logger.info("Fetched %d free models", len(free_models))
|
|
return free_models
|
|
except Exception as exc:
|
|
logger.warning("Failed to fetch free models: %s", exc)
|
|
finally:
|
|
if connector:
|
|
await connector.close()
|
|
|
|
if _free_models_cache:
|
|
return _free_models_cache
|
|
|
|
return fallback
|
|
|
|
|
|
async def _test_model(session: aiohttp.ClientSession, model: str) -> bool:
|
|
headers = {
|
|
"Authorization": f"Bearer {OPENROUTER_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": "reply OK"}],
|
|
"max_tokens": 5,
|
|
}
|
|
try:
|
|
async with session.post(OPENROUTER_URL, json=payload, headers=headers) as response:
|
|
if response.status == 200:
|
|
data = await response.json()
|
|
choices = data.get("choices", [])
|
|
return bool(choices and choices[0].get("message", {}).get("content"))
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
|
|
async def _update_working_models() -> None:
|
|
global _working_models_cache, _working_models_cache_time
|
|
|
|
free_models = await _fetch_free_models()
|
|
if not free_models:
|
|
return
|
|
|
|
connector = get_proxy_connector()
|
|
working: list[str] = []
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
for model in free_models:
|
|
if await _test_model(session, model):
|
|
working.append(model)
|
|
await asyncio.sleep(0.3)
|
|
finally:
|
|
if connector:
|
|
await connector.close()
|
|
|
|
_working_models_cache = working
|
|
_working_models_cache_time = time.time()
|
|
logger.info("Health check: %d/%d models working", len(working), len(free_models))
|
|
|
|
|
|
async def start_model_health_check() -> None:
|
|
if not AI_HEALTH_CHECK_ENABLED:
|
|
logger.info("Model health check is disabled")
|
|
return
|
|
|
|
await asyncio.sleep(30)
|
|
try:
|
|
await _update_working_models()
|
|
except Exception as exc:
|
|
logger.error("Initial health check error: %s", exc)
|
|
|
|
while True:
|
|
await asyncio.sleep(AI_HEALTH_CHECK_INTERVAL)
|
|
try:
|
|
await _update_working_models()
|
|
except Exception as exc:
|
|
logger.error("Health check error: %s", exc)
|
|
|
|
|
|
def _log_usage(source: str, model: str, data: dict, latency: float) -> None:
|
|
usage = data.get("usage")
|
|
if usage:
|
|
logger.info(
|
|
"AI %s | model=%s in_tok=%s out_tok=%s total_tok=%s latency=%.1fs",
|
|
source,
|
|
model,
|
|
usage.get("prompt_tokens", "?"),
|
|
usage.get("completion_tokens", "?"),
|
|
usage.get("total_tokens", "?"),
|
|
latency,
|
|
)
|
|
else:
|
|
logger.info("AI %s | model=%s latency=%.1fs", source, model, latency)
|
|
|
|
|
|
async def _try_openrouter(
|
|
session: aiohttp.ClientSession,
|
|
model: str,
|
|
headers: dict,
|
|
base_payload: dict,
|
|
raw: bool = False,
|
|
) -> tuple[str | None, bool]:
|
|
"""Returns (content, should_retry_later)."""
|
|
payload = {**base_payload, "model": model}
|
|
start = time.monotonic()
|
|
|
|
async with session.post(OPENROUTER_URL, json=payload, headers=headers) as response:
|
|
latency = time.monotonic() - start
|
|
|
|
if response.status == 429:
|
|
retry_after = response.headers.get("Retry-After")
|
|
logger.warning(
|
|
"OpenRouter rate limited | model=%s retry_after=%s",
|
|
model,
|
|
retry_after,
|
|
)
|
|
return None, True
|
|
|
|
if response.status != 200:
|
|
error_body = await response.text()
|
|
logger.warning(
|
|
"OpenRouter error | model=%s status=%s latency=%.1fs error=%s",
|
|
model,
|
|
response.status,
|
|
latency,
|
|
error_body[:200],
|
|
)
|
|
return None, False
|
|
|
|
data = await response.json()
|
|
choices = data.get("choices", [])
|
|
if not choices:
|
|
logger.warning("OpenRouter empty choices | model=%s latency=%.1fs", model, latency)
|
|
return None, False
|
|
|
|
content = choices[0].get("message", {}).get("content")
|
|
if not content:
|
|
logger.warning("OpenRouter empty content | model=%s latency=%.1fs", model, latency)
|
|
return None, False
|
|
|
|
_log_usage("OpenRouter", model, data, latency)
|
|
return content if raw else _md_to_html(content), False
|
|
|
|
|
|
async def _try_routerai(session: aiohttp.ClientSession, messages: list[dict]) -> str | None:
|
|
if not ROUTERAI_API_KEY:
|
|
logger.warning("RouterAI skipped | key not set")
|
|
return None
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {ROUTERAI_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload = {
|
|
"model": ROUTERAI_MODEL,
|
|
"messages": messages,
|
|
"max_tokens": 512,
|
|
}
|
|
|
|
start = time.monotonic()
|
|
|
|
async with session.post(ROUTERAI_URL, json=payload, headers=headers) as response:
|
|
latency = time.monotonic() - start
|
|
|
|
if response.status != 200:
|
|
error_body = await response.text()
|
|
logger.warning(
|
|
"RouterAI error | status=%s latency=%.1fs error=%s",
|
|
response.status,
|
|
latency,
|
|
error_body[:200],
|
|
)
|
|
return None
|
|
|
|
data = await response.json()
|
|
choices = data.get("choices", [])
|
|
if not choices:
|
|
logger.warning("RouterAI empty choices | latency=%.1fs", latency)
|
|
return None
|
|
|
|
content = choices[0].get("message", {}).get("content")
|
|
if not content:
|
|
logger.warning("RouterAI empty content | latency=%.1fs", latency)
|
|
return None
|
|
|
|
_log_usage("RouterAI", ROUTERAI_MODEL, data, latency)
|
|
return _md_to_html(content) + PAID_NOTICE
|
|
|
|
|
|
async def ask_ai_simple(prompt: str) -> str | None:
|
|
headers = {
|
|
"Authorization": f"Bearer {OPENROUTER_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
base_payload = {
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 512,
|
|
}
|
|
|
|
connector = get_proxy_connector()
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(
|
|
timeout=get_client_timeout(), connector=connector
|
|
) as session:
|
|
free_models = await _fetch_free_models()
|
|
for model in free_models[:5]:
|
|
result, _ = await _try_openrouter(session, model, headers, base_payload, raw=True)
|
|
if result:
|
|
# ask_ai_simple returns raw text for internal use (summaries)
|
|
return result
|
|
logger.warning("ask_ai_simple fallback fail | model=%s", model)
|
|
|
|
logger.warning("ask_ai_simple | all free models failed, trying RouterAI")
|
|
routerai_payload = {
|
|
"model": ROUTERAI_MODEL,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 512,
|
|
}
|
|
routerai_headers = {
|
|
"Authorization": f"Bearer {ROUTERAI_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
start = time.monotonic()
|
|
async with session.post(ROUTERAI_URL, json=routerai_payload, headers=routerai_headers) as response:
|
|
latency = time.monotonic() - start
|
|
if response.status == 200:
|
|
data = await response.json()
|
|
choices = data.get("choices", [])
|
|
if choices and choices[0].get("message", {}).get("content"):
|
|
_log_usage("ask_ai_simple (RouterAI)", ROUTERAI_MODEL, data, latency)
|
|
return choices[0]["message"]["content"]
|
|
except Exception as exc:
|
|
logger.error("ask_ai_simple error: %s", exc)
|
|
finally:
|
|
if connector:
|
|
await connector.close()
|
|
|
|
return None
|
|
|
|
|
|
async def ask_ai(
|
|
prompt: str,
|
|
context_messages: list[dict] | None = None,
|
|
status_callback=None,
|
|
extra_system_content: str | None = None,
|
|
) -> str:
|
|
system_content = AI_SYSTEM_PROMPT
|
|
if extra_system_content:
|
|
system_content += "\n\n" + extra_system_content
|
|
|
|
messages: list[dict] = [{"role": "system", "content": system_content}]
|
|
|
|
if context_messages:
|
|
messages.extend(context_messages)
|
|
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
or_headers = {
|
|
"Authorization": f"Bearer {OPENROUTER_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
"HTTP-Referer": "https://github.com/umb-bot",
|
|
"X-Title": "UMB Bot",
|
|
}
|
|
|
|
base_payload = {
|
|
"messages": messages,
|
|
"max_tokens": 512,
|
|
}
|
|
|
|
connector = get_proxy_connector()
|
|
|
|
waiting_messages = [
|
|
"Думаю...",
|
|
"Ой, надо ещё подумать...",
|
|
"Секундочку...",
|
|
"Ищу ответ...",
|
|
"Думаю...",
|
|
"Почти готово...",
|
|
"Переключаюсь на платный API...",
|
|
]
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(
|
|
timeout=get_client_timeout(), connector=connector
|
|
) as session:
|
|
free_models = await _fetch_free_models()
|
|
|
|
if _working_models_cache and (time.time() - _working_models_cache_time) < _working_models_cache_ttl:
|
|
models_to_try = [m for m in _working_models_cache if m in free_models]
|
|
if not models_to_try:
|
|
models_to_try = free_models
|
|
else:
|
|
models_to_try = free_models
|
|
|
|
logger.info("ask_ai | trying %d models", len(models_to_try))
|
|
|
|
for i, model in enumerate(models_to_try):
|
|
if status_callback and i > 0:
|
|
wait_idx = min(i, len(waiting_messages) - 1)
|
|
await status_callback(waiting_messages[wait_idx])
|
|
|
|
result, _ = await _try_openrouter(session, model, or_headers, base_payload)
|
|
if result:
|
|
return result
|
|
|
|
logger.info("ask_ai | model %s failed, trying next", model)
|
|
|
|
logger.warning("ask_ai | switching to RouterAI")
|
|
|
|
if status_callback:
|
|
await status_callback(waiting_messages[-1])
|
|
|
|
paid_result = await _try_routerai(session, messages)
|
|
if paid_result:
|
|
return paid_result
|
|
|
|
logger.warning("ask_ai | RouterAI failed too")
|
|
|
|
logger.error("ask_ai | all models failed")
|
|
return "Извини, ни одна модель не смогла ответить. Попробуй позже."
|
|
except aiohttp.ClientError as exc:
|
|
logger.error("AI request network error: %s", exc)
|
|
return "Не удалось связаться с AI сервисом. Проверь соединение."
|
|
except asyncio.TimeoutError:
|
|
logger.error("ask_ai | timeout after all models exhausted")
|
|
return "Сервер AI не ответил вовремя. Попробуй позже."
|
|
except Exception:
|
|
logger.exception("ask_ai | unexpected error")
|
|
return "Произошла ошибка при обработке запроса."
|
|
finally:
|
|
if connector:
|
|
await connector.close()
|