mirror of
https://github.com/MuRuLOSE/limoka.git
synced 2026-06-16 14:34:17 +02:00
389 lines
14 KiB
Python
389 lines
14 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
|
||
# ---------------------------------------------------------------------------------
|
||
|
||
import html
|
||
import httpx
|
||
import asyncio
|
||
import logging
|
||
|
||
from openai import AsyncOpenAI
|
||
from typing import Optional, Any
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timezone, timedelta
|
||
|
||
from telethon.errors import MessageNotModifiedError
|
||
|
||
from .. import loader, utils
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@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
|
||
|
||
|
||
@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": "<emoji document_id=5197252827247841976>🐣</emoji> <b>Обрабатываю запрос в Cocoon...</b>",
|
||
"thinking": (
|
||
"<emoji document_id=5197252827247841976>🐣</emoji> <b>Думаю...</b>\n\n"
|
||
"<blockquote>{thoughts}…</blockquote>"
|
||
),
|
||
"answer": (
|
||
"<emoji document_id=5456217626957091223>🌘</emoji> <b>Вопрос:</b> {question}\n\n"
|
||
"<emoji document_id=5197252827247841976>🐣</emoji> <b>Размышления:</b>\n"
|
||
"<blockquote expandable>{thoughts}</blockquote>\n\n"
|
||
"<emoji document_id=5208521532942358129>🥚</emoji> {answer}\n\n"
|
||
"<emoji document_id=5458567764341985638>🚀</emoji> <b>Модель</b>: <code>{model}</code>"
|
||
),
|
||
"usage": (
|
||
"<b><emoji document_id=5208521532942358129>🥚</emoji> Cocoon API\n\n"
|
||
"<emoji document_id=5458805877328875335>💡</emoji> Использовано:\n"
|
||
"</b><i>• {current}/{total} ({percent}% осталось)</i><b>\n\n"
|
||
"<emoji document_id=5456591761558245861>⏳</emoji> Лимит сбросится через {days} день(-ей).</b>"
|
||
),
|
||
}
|
||
|
||
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).",
|
||
),
|
||
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):
|
||
try:
|
||
if len(text) > 4096:
|
||
text = text[:4090] + "..."
|
||
return await utils.answer(message, text)
|
||
except MessageNotModifiedError:
|
||
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),
|
||
)
|
||
|
||
@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):
|
||
"""Задать вопрос к ИИ"""
|
||
|
||
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"])
|
||
|
||
message = await utils.answer(message, self.strings["sending_request_to_cocoon"])
|
||
|
||
self._ensure_client()
|
||
|
||
client = AsyncOpenAI(api_key=self.config["token"], base_url=self.api_url)
|
||
|
||
system_prompt = (self.config.get("system_prompt") or "").strip()
|
||
|
||
messages = []
|
||
if system_prompt:
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
messages.append({"role": self.config["role"] or "user", "content": q})
|
||
|
||
try:
|
||
response = await client.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) >= 150:
|
||
response_text += chunk_buffer
|
||
chunk_buffer = ""
|
||
|
||
thoughts = (
|
||
response_text.replace("<think>", "")
|
||
.replace("</think>", "")
|
||
.strip()
|
||
)
|
||
|
||
if "</think>" in response_text:
|
||
after_think = response_text.split("</think>", 1)[-1].strip()
|
||
await self._answer(
|
||
message,
|
||
self.strings["answer"].format(
|
||
thoughts=thoughts[:300],
|
||
question=q,
|
||
answer=_escape_text(after_think) + "…",
|
||
model=self.config["model"],
|
||
),
|
||
)
|
||
else:
|
||
thinking_text = (
|
||
response_text.replace("<think>", "")
|
||
.replace("</think>", "")
|
||
.strip()
|
||
)
|
||
await self._answer(
|
||
message,
|
||
self.strings["thinking"].format(
|
||
thoughts=_escape_text(thinking_text)
|
||
),
|
||
)
|
||
|
||
await asyncio.sleep(2)
|
||
|
||
if chunk_buffer:
|
||
response_text += chunk_buffer
|
||
|
||
if "</think>" in response_text:
|
||
after_think = response_text.split("</think>", 1)[-1].strip()
|
||
else:
|
||
after_think = (
|
||
response_text.replace("<think>", "").replace("</think>", "").strip()
|
||
)
|
||
|
||
await self._answer(
|
||
message,
|
||
self.strings["answer"].format(
|
||
thoughts=thoughts,
|
||
question=q,
|
||
answer=_escape_text(after_think),
|
||
model=self.config["model"],
|
||
),
|
||
)
|
||
|
||
except httpx.RemoteProtocolError:
|
||
return await self._answer(message, self.strings["try_again"])
|