mirror of
https://github.com/MuRuLOSE/limoka.git
synced 2026-06-16 06:24:18 +02:00
554 lines
19 KiB
Python
554 lines
19 KiB
Python
# █▀▀ ▄▀█ █▀▄▀█ █▀█ █▀▄ █▀
|
||
# █▀░ █▀█ █░▀░█ █▄█ █▄▀ ▄█
|
||
|
||
# https://t.me/famods
|
||
|
||
# 🔒 Licensed under the GNU AGPLv3
|
||
# 🌐 https://www.gnu.org/licenses/agpl-3.0.html
|
||
|
||
# ---------------------------------------------------------------------------------
|
||
# Name: Cocoon [BETA]
|
||
# Description: Взаимодействие с Cocoon от HikkaHost
|
||
# meta developer: @FAmods & @vsecoder_m
|
||
# meta banner: https://github.com/FajoX1/FAmods/blob/main/assets/banners/cocoon.png?raw=true
|
||
# requires: openai httpx aiohttp bs4 markdown
|
||
# ---------------------------------------------------------------------------------
|
||
|
||
import re
|
||
import html
|
||
import uuid
|
||
import httpx
|
||
import asyncio
|
||
import logging
|
||
import markdown
|
||
|
||
from openai import AsyncOpenAI
|
||
from typing import Optional, Any
|
||
from dataclasses import dataclass
|
||
from bs4 import BeautifulSoup, NavigableString
|
||
from datetime import datetime, timezone, timedelta
|
||
|
||
from telethon.tl.types import User
|
||
from telethon.errors import MessageNotModifiedError
|
||
|
||
from aiogram.exceptions import TelegramBadRequest
|
||
|
||
from .. import loader, utils
|
||
from ..inline.types import InlineCall
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
TG_ALLOWED = {
|
||
"b",
|
||
"strong",
|
||
"i",
|
||
"em",
|
||
"u",
|
||
"ins",
|
||
"s",
|
||
"strike",
|
||
"del",
|
||
"a",
|
||
"code",
|
||
"pre",
|
||
"blockquote",
|
||
"emoji",
|
||
"tg-emoji",
|
||
}
|
||
TAG_MAP = {"strong": "b", "em": "i", "del": "s", "strike": "s", "ins": "u"}
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Usage:
|
||
spent_nano: int
|
||
spent_ton: str
|
||
tokens_spent: int
|
||
free_tokens_remaining: int
|
||
free_tokens_reset_at: Optional[int]
|
||
updated_at: Optional[int]
|
||
|
||
|
||
def _now_utc() -> datetime:
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def _safe_int(v: Any, default: int = 0) -> int:
|
||
try:
|
||
return int(v)
|
||
except Exception:
|
||
return default
|
||
|
||
|
||
def _escape_text(s: str) -> str:
|
||
return html.escape(s or "", quote=False)
|
||
|
||
|
||
def _days_until_reset(reset_ts: Optional[int], cycle_days: int = 30) -> int:
|
||
"""
|
||
Cocoon может отдавать free_tokens_reset_at как момент сброса текущего периода,
|
||
который уже мог произойти сегодня. Тогда следующий сброс = reset_ts + cycle_days.
|
||
"""
|
||
if not reset_ts:
|
||
return cycle_days
|
||
|
||
try:
|
||
now = _now_utc()
|
||
target = datetime.fromtimestamp(int(reset_ts), tz=timezone.utc)
|
||
|
||
if target <= now:
|
||
target = target + timedelta(days=cycle_days)
|
||
|
||
delta = target - now
|
||
if delta.total_seconds() <= 0:
|
||
return 0
|
||
|
||
return max(delta.days + (1 if delta.seconds > 0 else 0), 0)
|
||
except Exception:
|
||
return cycle_days
|
||
|
||
|
||
def _normalize_reset_ts(reset_at: Optional[int]) -> Optional[int]:
|
||
if not reset_at:
|
||
return None
|
||
|
||
reset_at = _safe_int(reset_at, 0)
|
||
if reset_at <= 0:
|
||
return None
|
||
|
||
return reset_at
|
||
|
||
|
||
def _format_compact(n: int) -> str:
|
||
n = int(n)
|
||
|
||
if n >= 1_000_000_000:
|
||
return f"{n / 1_000_000_000:.1f}".rstrip("0").rstrip(".") + "b"
|
||
if n >= 1_000_000:
|
||
return f"{n / 1_000_000:.1f}".rstrip("0").rstrip(".") + "m"
|
||
if n >= 1_000:
|
||
return f"{n / 1_000:.1f}".rstrip("0").rstrip(".") + "k"
|
||
return str(n)
|
||
|
||
|
||
def _percent_remaining(spent: int, total: int) -> float:
|
||
if total <= 0:
|
||
return 0.0
|
||
remaining = max(total - spent, 0)
|
||
return (remaining / total) * 100.0
|
||
|
||
|
||
def md_to_tg_html(text: str) -> str:
|
||
if not text:
|
||
return ""
|
||
|
||
raw_html = markdown.markdown(text, extensions=["fenced_code", "tables", "nl2br"])
|
||
soup = BeautifulSoup(raw_html, "html.parser")
|
||
|
||
def stringify(node, lang=None):
|
||
res = ""
|
||
|
||
for child in node.children:
|
||
if isinstance(child, NavigableString):
|
||
res += html.escape(str(child))
|
||
|
||
elif child.name:
|
||
tag_name = child.name
|
||
|
||
if tag_name in ["h1", "h2", "h3", "h4", "h5", "h6"]:
|
||
content = stringify(child)
|
||
res += f"<b>{content}</b>\n\n"
|
||
elif tag_name == "p":
|
||
res += stringify(child) + "\n\n"
|
||
elif tag_name == "br":
|
||
res += "\n"
|
||
elif tag_name == "li":
|
||
res += f"• {stringify(child)}\n"
|
||
elif tag_name in ["ul", "ol"]:
|
||
res += stringify(child) + "\n"
|
||
elif tag_name == "tr":
|
||
res += "| " + stringify(child) + "\n"
|
||
elif tag_name in ["td", "th"]:
|
||
res += stringify(child) + " | "
|
||
|
||
else:
|
||
target_tag = TAG_MAP.get(tag_name, tag_name)
|
||
|
||
if target_tag in TG_ALLOWED:
|
||
inner_html = stringify(child)
|
||
|
||
if not inner_html.strip() and target_tag not in ["code", "pre"]:
|
||
res += inner_html
|
||
continue
|
||
|
||
if target_tag == "a":
|
||
href = child.get("href", "")
|
||
if href:
|
||
res += f'<a href="{html.escape(href)}">{inner_html}</a>'
|
||
else:
|
||
res += inner_html
|
||
elif target_tag == "code":
|
||
cls = child.get("class", [])
|
||
if cls and cls[0].startswith("language-"):
|
||
res += f'<code class="{cls[0]}">{inner_html}</code>'
|
||
else:
|
||
res += f"<code>{inner_html}</code>"
|
||
elif target_tag == "pre":
|
||
res += f"<pre>{inner_html}</pre>"
|
||
else:
|
||
res += f"<{target_tag}>{inner_html}</{target_tag}>"
|
||
else:
|
||
res += stringify(child)
|
||
return res
|
||
|
||
final_text = stringify(soup)
|
||
|
||
final_text = re.sub(r"\n{3,}", "\n\n", final_text)
|
||
return final_text.strip()
|
||
|
||
|
||
def repair_html_tags(html_chunk: str) -> str:
|
||
if not html_chunk:
|
||
return ""
|
||
|
||
newline_placeholder = f"MARKER_{uuid.uuid4().hex}"
|
||
|
||
protected_html = html_chunk.replace("\n", newline_placeholder)
|
||
|
||
soup = BeautifulSoup(protected_html, "html.parser")
|
||
|
||
repaired_html = soup.decode_contents(formatter=None)
|
||
|
||
final_html = repaired_html.replace(newline_placeholder, "\n")
|
||
|
||
return final_html
|
||
|
||
|
||
@loader.tds
|
||
class Cocoon(loader.Module):
|
||
"""Взаимодействие с Cocoon от HikkaHost"""
|
||
|
||
strings = {
|
||
"name": "Cocoon [BETA]",
|
||
"try_again": "<b><emoji document_id=5456307331644037599>❌</emoji> <b>Что-то пошло не так. Попробуйте снова.</b>",
|
||
"no_args": "<b><emoji document_id=5456307331644037599>❌</emoji> <b>Нужно </b><code>{}{} {}</code></b>",
|
||
"no_token": (
|
||
"<b><emoji document_id=5456307331644037599>❌</emoji> <b>Нету токена! Вставь его в <code>{}cfg cocoon</code>\n\n"
|
||
"<emoji document_id=5456672880605565619>🌘</emoji> Получить токен: @hikkahost_bot → <emoji document_id=5208521532942358129>🥚</emoji> Cocoon</b>"
|
||
),
|
||
"invalid_token_or_no_sub": (
|
||
"<b><emoji document_id=5456307331644037599>❌</emoji>Неверный токен или у вас нет подписки <emoji document_id=5188377234380954537>🌘</emoji> HikkaHost.</b>\n\n"
|
||
"<emoji document_id=5456672880605565619>🌘</emoji> Получить токен: @hikkahost_bot → <emoji document_id=5208521532942358129>🥚</emoji> Cocoon</b>"
|
||
),
|
||
"sending_request_to_cocoon": "<tg-emoji emoji-id=5197252827247841976>🐣</tg-emoji> <b>Обрабатываю запрос в Cocoon...</b>",
|
||
"thinking": (
|
||
"<tg-emoji emoji-id=5197252827247841976>🐣</tg-emoji> <b>Думаю...</b>\n\n"
|
||
"<blockquote>{thoughts}…</blockquote>"
|
||
),
|
||
"answer": (
|
||
"<tg-emoji emoji-id=5456217626957091223>🌘</tg-emoji> <b>Вопрос:</b> {question}\n\n"
|
||
"<tg-emoji emoji-id=5197252827247841976>🐣</tg-emoji> <b>Размышления:</b>\n"
|
||
"<blockquote expandable>{thoughts}</blockquote>\n\n"
|
||
"<tg-emoji emoji-id=5208521532942358129>🥚</tg-emoji> {answer}\n\n"
|
||
"<tg-emoji emoji-id=5458567764341985638>🚀</tg-emoji> <b>Модель</b>: <code>{model}</code>"
|
||
),
|
||
"usage": (
|
||
"<b><tg-emoji emoji-id=5208521532942358129>🥚</tg-emoji> Cocoon API\n\n"
|
||
"<tg-emoji emoji-id=5458805877328875335>💡</tg-emoji> Использовано:\n"
|
||
"</b><i>• {current}/{total} ({percent}% осталось)</i><b>\n\n"
|
||
"<tg-emoji emoji-id=5456591761558245861>⏳</tg-emoji> Лимит сбросится через {days} день(-ей).</b>"
|
||
),
|
||
"again_kb": "🔄 Сгенерировать ещё раз",
|
||
}
|
||
|
||
def __init__(self):
|
||
self.config = loader.ModuleConfig(
|
||
loader.ConfigValue(
|
||
"token",
|
||
None,
|
||
lambda: "Токен HikkaHost API. Получить токен: @hikkahost_bot -> 🥚 Cocoon",
|
||
validator=loader.validators.Hidden(loader.validators.String()),
|
||
),
|
||
loader.ConfigValue(
|
||
"model",
|
||
"Qwen/Qwen3-32B",
|
||
lambda: "Модель ИИ. Список: https://cocoon.hikka.host/v1/models",
|
||
),
|
||
loader.ConfigValue(
|
||
"role",
|
||
"user",
|
||
lambda: "Роль user-сообщения (обычно user).",
|
||
validator=loader.validators.Choice(
|
||
["system", "developer", "user", "assistant", "tool"]
|
||
),
|
||
),
|
||
loader.ConfigValue(
|
||
"system_prompt",
|
||
"",
|
||
lambda: "System prompt (инструкция для модели, role=system).",
|
||
),
|
||
loader.ConfigValue(
|
||
"max_tokens",
|
||
3900,
|
||
lambda: "Максимальное количество токенов для ответа модели.",
|
||
validator=loader.validators.Integer(minimum=1),
|
||
),
|
||
loader.ConfigValue(
|
||
"temperature",
|
||
0.2,
|
||
lambda: "Температура генерации (0.0–1.0).",
|
||
),
|
||
)
|
||
|
||
async def client_ready(self, client, db):
|
||
self.db = db
|
||
self._client = client
|
||
|
||
self.api_url = "https://cocoon.hikka.host/v1"
|
||
self._openai: Optional[AsyncOpenAI] = None
|
||
|
||
def _rebuild_openai_client(self) -> None:
|
||
token = self.config.get("token") or ""
|
||
self._openai = AsyncOpenAI(
|
||
api_key=token, base_url=self.api_url, timeout=60.0, max_retries=2
|
||
)
|
||
|
||
def _ensure_client(self) -> None:
|
||
if not self._openai or (
|
||
self._openai.api_key != (self.config.get("token") or "")
|
||
):
|
||
self._rebuild_openai_client()
|
||
|
||
async def _answer(self, message, text, *args, **kwargs):
|
||
try:
|
||
if len(text) > 4096:
|
||
text = text[:4090] + "..."
|
||
|
||
return await utils.answer(message, repair_html_tags(text), *args, **kwargs)
|
||
except (MessageNotModifiedError, TelegramBadRequest):
|
||
return message
|
||
|
||
async def _fetch_usage(self) -> Optional[Usage]:
|
||
token = self.config.get("token")
|
||
if not token:
|
||
return None
|
||
|
||
headers = {
|
||
"accept": "application/json",
|
||
"X-API-Key": token,
|
||
}
|
||
|
||
try:
|
||
async with httpx.AsyncClient(timeout=20.0) as client:
|
||
r = await client.get(f"{self.api_url}/usage", headers=headers)
|
||
r.raise_for_status()
|
||
data = r.json()
|
||
except Exception as e:
|
||
logger.exception("Usage request failed: %s", e)
|
||
return None
|
||
|
||
if isinstance(data, dict) and data.get("detail") == "API key not recognized":
|
||
return None
|
||
|
||
if not isinstance(data, dict):
|
||
return None
|
||
|
||
return Usage(
|
||
spent_nano=_safe_int(data.get("spent_nano"), 0),
|
||
spent_ton=str(data.get("spent_ton", "0")),
|
||
tokens_spent=_safe_int(data.get("tokens_spent"), 0),
|
||
free_tokens_remaining=_safe_int(data.get("free_tokens_remaining"), 0),
|
||
free_tokens_reset_at=(
|
||
_safe_int(data.get("free_tokens_reset_at"), 0) or None
|
||
),
|
||
updated_at=(_safe_int(data.get("updated_at"), 0) or None),
|
||
)
|
||
|
||
async def _regenerate(self, call: InlineCall, arg1, arg2):
|
||
await self.cocoon(arg1, inline_message=arg2)
|
||
|
||
@loader.command()
|
||
async def ccusage(self, message):
|
||
"""Статистика использования Cocoon"""
|
||
|
||
if not self.config.get("token"):
|
||
return await self._answer(
|
||
message, self.strings["no_token"].format(self.get_prefix())
|
||
)
|
||
|
||
usage = await self._fetch_usage()
|
||
if not usage:
|
||
return await self._answer(message, self.strings["invalid_token_or_no_sub"])
|
||
|
||
reset_ts = _normalize_reset_ts(usage.free_tokens_reset_at)
|
||
days = _days_until_reset(reset_ts, cycle_days=30)
|
||
|
||
spent = usage.tokens_spent
|
||
total = spent + usage.free_tokens_remaining
|
||
if total <= 0:
|
||
total = 1_000_000
|
||
|
||
percent = _percent_remaining(spent, total)
|
||
percent_fmt = f"{percent:.1f}".rstrip("0").rstrip(".")
|
||
|
||
return await self._answer(
|
||
message,
|
||
self.strings["usage"].format(
|
||
current=_format_compact(spent),
|
||
total=_format_compact(total),
|
||
percent=percent_fmt,
|
||
days=days,
|
||
),
|
||
)
|
||
|
||
@loader.command()
|
||
async def cocoon(self, message, inline_message=None):
|
||
"""Задать вопрос к ИИ (поддерживает ответ на сообщение)"""
|
||
|
||
q = utils.get_args_raw(message)
|
||
if not q:
|
||
return await utils.answer(
|
||
message,
|
||
self.strings["no_args"].format(self.get_prefix(), "cocoon", "[вопрос]"),
|
||
)
|
||
|
||
if not self.config["token"]:
|
||
return await utils.answer(
|
||
message, self.strings["no_token"].format(self.get_prefix())
|
||
)
|
||
|
||
usage = await self._fetch_usage()
|
||
if not usage:
|
||
return await utils.answer(message, self.strings["invalid_token_or_no_sub"])
|
||
|
||
user_message = message
|
||
|
||
if not inline_message:
|
||
message = await self.inline.form(
|
||
text="...",
|
||
message=message,
|
||
force_me=False,
|
||
)
|
||
else:
|
||
message = inline_message
|
||
|
||
await utils.answer(message, self.strings["sending_request_to_cocoon"])
|
||
|
||
self._ensure_client()
|
||
|
||
system_prompt = (self.config.get("system_prompt") or "").strip()
|
||
|
||
messages = []
|
||
if system_prompt:
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
|
||
if user_message.reply_to:
|
||
reply = await user_message.get_reply_message()
|
||
|
||
entity_id = None
|
||
|
||
if hasattr(reply, "from_id") and reply.from_id:
|
||
if hasattr(reply.from_id, "user_id") and reply.from_id.user_id:
|
||
entity_id = reply.from_id.user_id
|
||
elif hasattr(reply.from_id, "channel_id") and reply.from_id.channel_id:
|
||
entity_id = reply.from_id.channel_id
|
||
|
||
if entity_id is None:
|
||
if hasattr(reply.peer_id, "user_id") and reply.peer_id.user_id:
|
||
entity_id = reply.peer_id.user_id
|
||
else:
|
||
entity_id = reply.peer_id.channel_id
|
||
|
||
entity = await self.client.get_entity(entity_id)
|
||
|
||
date = reply.date.strftime("%H:%M %d.%m.%Y UTC")
|
||
|
||
messages.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
(
|
||
f"{entity.first_name} {entity.last_name or ''} (user id: {entity.id}) "
|
||
if isinstance(entity, User)
|
||
else f"Channel {entity.title} (channel id: {entity.id}) "
|
||
)
|
||
+ f"msg id: {reply.id}, date: {date}: "
|
||
+ reply.message
|
||
),
|
||
}
|
||
)
|
||
|
||
messages.append({"role": self.config.get("role", "user"), "content": q})
|
||
|
||
try:
|
||
response = await self._openai.chat.completions.create(
|
||
messages=messages,
|
||
stream=True,
|
||
max_tokens=self.config.get("max_tokens", 3900),
|
||
model=self.config.get("model", "Qwen/Qwen3-32B"),
|
||
temperature=self.config.get("temperature", 0.2),
|
||
)
|
||
|
||
response_text = ""
|
||
chunk_buffer = ""
|
||
|
||
async for chunk in response:
|
||
if chunk.choices and chunk.choices[0].delta.content:
|
||
chunk_buffer += chunk.choices[0].delta.content
|
||
|
||
if len(chunk_buffer) >= 100:
|
||
response_text += chunk_buffer
|
||
chunk_buffer = ""
|
||
|
||
thoughts = response_text.split("</think>", 1)[0].replace(
|
||
"<think>", ""
|
||
)
|
||
|
||
if "</think>" in response_text:
|
||
after_think = response_text.split("</think>", 1)[1].strip()
|
||
await self._answer(
|
||
message,
|
||
self.strings["answer"].format(
|
||
thoughts=thoughts[:500],
|
||
question=q,
|
||
answer=md_to_tg_html(_escape_text(after_think) + "…"),
|
||
model=self.config["model"],
|
||
),
|
||
)
|
||
else:
|
||
await self._answer(
|
||
message,
|
||
self.strings["thinking"].format(
|
||
thoughts=md_to_tg_html(_escape_text(thoughts) + "…")
|
||
),
|
||
)
|
||
|
||
await asyncio.sleep(0.2)
|
||
|
||
if chunk_buffer:
|
||
response_text += chunk_buffer
|
||
|
||
responses_data = response_text.split("</think>", 1)
|
||
thoughts = responses_data[0].strip().replace("<think>", "")
|
||
after_think = responses_data[1].strip()
|
||
|
||
await self._answer(
|
||
message,
|
||
self.strings["answer"].format(
|
||
question=q,
|
||
thoughts=md_to_tg_html(_escape_text(thoughts[:500])),
|
||
answer=md_to_tg_html(_escape_text(after_think)),
|
||
model=self.config["model"],
|
||
),
|
||
reply_markup=[
|
||
{
|
||
"text": self.strings["again_kb"],
|
||
"callback": self._regenerate,
|
||
"args": [user_message, message],
|
||
}
|
||
],
|
||
)
|
||
|
||
except httpx.RemoteProtocolError:
|
||
return await self._answer(message, self.strings["try_again"])
|