Initial commit: UMB Telegram Bot
This commit is contained in:
Галингер Р.С.
@@ -0,0 +1,401 @@
|
||||
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 "Произошла ошибка при обработке запроса."
|
||||
Reference in New Issue
Block a user