402 lines
14 KiB
Python
402 lines
14 KiB
Python
import asyncio
|
|
import logging
|
|
import re
|
|
import time
|
|
|
|
import aiohttp
|
|
from aiohttp_socks import ProxyConnector, ProxyType
|
|
|
|
from config import (
|
|
OPENROUTER_API_KEY,
|
|
AI_SYSTEM_PROMPT,
|
|
PROXY_ENABLED,
|
|
PROXY_URL,
|
|
ROUTERAI_API_KEY,
|
|
ROUTERAI_BASE_URL,
|
|
ROUTERAI_MODEL,
|
|
)
|
|
|
|
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"
|
|
|
|
_request_timeout = aiohttp.ClientTimeout(total=60, sock_connect=15, sock_read=30)
|
|
|
|
|
|
def get_client_timeout(total: int = 60) -> aiohttp.ClientTimeout:
|
|
return aiohttp.ClientTimeout(total=total, sock_connect=15, sock_read=30)
|
|
|
|
|
|
_free_models_cache = []
|
|
_free_models_cache_time = 0
|
|
_free_models_cache_ttl = 3600
|
|
|
|
_working_models_cache = []
|
|
_working_models_cache_time = 0
|
|
_working_models_cache_ttl = 600
|
|
|
|
|
|
def _get_connector():
|
|
if PROXY_ENABLED and PROXY_URL:
|
|
parsed = PROXY_URL.replace("socks5://", "").replace("socks5h://", "")
|
|
if "@" in parsed:
|
|
auth, host_port = parsed.split("@", 1)
|
|
username, password = auth.split(":", 1)
|
|
else:
|
|
username = None
|
|
password = None
|
|
host_port = parsed
|
|
|
|
host, port = host_port.rsplit(":", 1)
|
|
port = int(port)
|
|
|
|
return ProxyConnector(
|
|
proxy_type=ProxyType.SOCKS5,
|
|
host=host,
|
|
port=port,
|
|
username=username,
|
|
password=password,
|
|
)
|
|
return None
|
|
|
|
|
|
def _md_to_html(text: str) -> str:
|
|
text = re.sub(r"```(\w*)\n(.*?)```", r"<pre>\2</pre>", 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",
|
|
}
|
|
|
|
try:
|
|
connector = _get_connector()
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
async with session.get("https://openrouter.ai/api/v1/models", headers=headers) as response:
|
|
if response.status == 200:
|
|
data = await response.json()
|
|
models = data.get("data", [])
|
|
|
|
free_models = []
|
|
for model in models:
|
|
model_id = model.get("id", "")
|
|
if model_id.endswith(":free"):
|
|
free_models.append(model_id)
|
|
|
|
if free_models:
|
|
_free_models_cache = free_models
|
|
_free_models_cache_time = now
|
|
logger.info(f"Fetched {len(free_models)} free models")
|
|
return free_models
|
|
except Exception as e:
|
|
logger.warning(f"Failed to fetch free models: {e}")
|
|
|
|
if _free_models_cache:
|
|
return _free_models_cache
|
|
|
|
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",
|
|
]
|
|
return fallback
|
|
|
|
|
|
async def _test_model(session, 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]["message"].get("content"))
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
|
|
async def _update_working_models():
|
|
global _working_models_cache, _working_models_cache_time
|
|
|
|
free_models = await _fetch_free_models()
|
|
if not free_models:
|
|
return
|
|
|
|
connector = _get_connector()
|
|
working = []
|
|
|
|
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)
|
|
|
|
_working_models_cache = working
|
|
_working_models_cache_time = time.time()
|
|
logger.info(f"Health check: {len(working)}/{len(free_models)} models working")
|
|
|
|
|
|
async def start_model_health_check():
|
|
await asyncio.sleep(30)
|
|
try:
|
|
await _update_working_models()
|
|
except Exception as e:
|
|
logger.error(f"Initial health check error: {e}")
|
|
|
|
while True:
|
|
await asyncio.sleep(600)
|
|
try:
|
|
await _update_working_models()
|
|
except Exception as e:
|
|
logger.error(f"Health check error: {e}")
|
|
|
|
|
|
def _log_usage(source: str, model: str, data: dict, latency: float):
|
|
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, model: str, messages: list[dict], headers: dict, payload: dict) -> str | None:
|
|
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 != 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
|
|
|
|
data = await response.json()
|
|
choices = data.get("choices", [])
|
|
if not choices:
|
|
logger.warning("OpenRouter empty choices | model=%s latency=%.1fs", model, latency)
|
|
return None
|
|
|
|
content = choices[0]["message"].get("content")
|
|
if not content:
|
|
logger.warning("OpenRouter empty content | model=%s latency=%.1fs", model, latency)
|
|
return None
|
|
|
|
_log_usage("OpenRouter", model, data, latency)
|
|
return _md_to_html(content)
|
|
|
|
|
|
async def _try_routerai(session, 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]["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",
|
|
}
|
|
|
|
payload = {
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 512,
|
|
}
|
|
|
|
connector = _get_connector()
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(timeout=_request_timeout, connector=connector) as session:
|
|
free_models = await _fetch_free_models()
|
|
for model in free_models[:5]:
|
|
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 == 200:
|
|
data = await response.json()
|
|
choices = data.get("choices", [])
|
|
if choices and choices[0]["message"].get("content"):
|
|
_log_usage("ask_ai_simple", model, data, latency)
|
|
return choices[0]["message"]["content"]
|
|
else:
|
|
logger.warning("ask_ai_simple fallback fail | model=%s status=%s latency=%.1fs", model, response.status, latency)
|
|
|
|
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]["message"].get("content"):
|
|
_log_usage("ask_ai_simple (RouterAI)", ROUTERAI_MODEL, data, latency)
|
|
return choices[0]["message"]["content"]
|
|
except Exception as e:
|
|
logger.error("ask_ai_simple error: %s", e)
|
|
|
|
return None
|
|
|
|
|
|
async def ask_ai(prompt: str, context_messages: list[dict] | None = None, status_callback=None) -> str:
|
|
messages = [{"role": "system", "content": AI_SYSTEM_PROMPT}]
|
|
|
|
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",
|
|
}
|
|
|
|
payload = {
|
|
"messages": messages,
|
|
"max_tokens": 512,
|
|
}
|
|
|
|
timeout = aiohttp.ClientTimeout(total=60, sock_connect=15, sock_read=30)
|
|
connector = _get_connector()
|
|
|
|
waiting_messages = [
|
|
"Думаю...",
|
|
"Ой, надо ещё подумать...",
|
|
"Секундочку...",
|
|
"Ищу ответ...",
|
|
"Думаю...",
|
|
"Почти готово...",
|
|
"Переключаюсь на платный API...",
|
|
]
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
|
|
free_models = await _fetch_free_models()
|
|
|
|
if _working_models_cache:
|
|
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, messages, or_headers, 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 e:
|
|
logger.error("AI request network error: %s", e)
|
|
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 "Произошла ошибка при обработке запроса."
|