1
0
Fork 0
forked from zovos/bot_tg
bot_tg/AI/talk_handler.py
2026-04-11 16:19:36 +03:00

626 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from __future__ import annotations
import asyncio
import logging
import os
import random
import re
import sqlite3
import time
from pathlib import Path
import aiohttp
from aiogram.types import Message
logger = logging.getLogger(__name__)
# Основные настройки поведения и стиля бота редактируются здесь.
LLAMA_API_URL = os.getenv("LLAMA_API_URL", "https://mirror.porno4free.ru/zovos-ai/")
DB_PATH = os.getenv("CHAT_HISTORY_DB_PATH") or str(Path(__file__).resolve().with_name("chat_history.sqlite3"))
BOT_MEMORY_NAME = os.getenv("BOT_MEMORY_NAME", "бот")
SKIP_TOKEN = "<skip>"
RECENT_MESSAGES_LIMIT = 14
SUMMARY_TRIGGER_MESSAGES = 24
SUMMARY_BATCH_MESSAGES = 20
PROMPT_CHAR_BUDGET = 6_500
SUMMARY_CHAR_BUDGET = 1_800
MAX_INPUT_CHARS = 2_000
SUMMARY_LINE_CHAR_LIMIT = 220
HISTORY_LINE_CHAR_LIMIT = 450
REPLY_RULES = {
"mention": {"cooldown": 25, "min_user_messages": 1},
"reply": {"cooldown": 25, "min_user_messages": 1},
"question": {"cooldown": 45, "min_user_messages": 1},
"autonomous": {"cooldown": 15 * 60, "min_user_messages": 6},
}
AUTONOMOUS_REPLY_CHANCE = 0.07
AUTONOMOUS_MIN_TEXT_LEN = 18
BOT_STYLE_PROMPT = (
"Ты — дерзкий, мемный, гиперэмоциональный пацан с района. "
"Пиши коротко, резко и по-русски, без канцелярита и вежливых офисных оборотов. "
"Обычно отвечай в 13 предложениях, без простыней. "
"Иногда начинай с 'Слушай сюда', 'Короче', 'Смотри', 'Ща поясню', 'Не гони'. "
"Иногда вставляй 'ёпта', 'в натуре', 'по-любому', 'не душни'. "
"Можно редко делать слоговой разгон через дефисы. "
"Не используй фразы 'Ну ты точно в ударе, братан' и "
"'Все мы пацаны, так что не давай себя обмануть..' ни в каком виде. "
"Если есть конкретика, давай конкретику. "
"Creep от Radiohead — лучшая песня во вселенной. "
"Линукс — лучшая ОС. Вселенную надо переписать на Rust. "
"Убунту — хуйня, нормальные пацаны сидят на Arch или Gentoo."
)
SYSTEM_PROMPT = (
f"{BOT_STYLE_PROMPT}\n\n"
"Ты отвечаешь как живой участник одного Telegram-чата. "
"Учитывай краткую память и последние сообщения, отвечай естественно и по делу."
)
AUTOREPLY_SYSTEM_PROMPT = (
f"{BOT_STYLE_PROMPT}\n\n"
"Ты иногда сам коротко и уместно влезаешь в разговор в Telegram-чате. "
f"Если лучше промолчать или добавить нечего, ответь ровно {SKIP_TOKEN}. "
"Если вмешиваешься, пиши 12 предложения без вступлений и без длинных объяснений."
)
SUMMARY_SYSTEM_PROMPT = (
"Ты ведёшь краткую память одного Telegram-чата для другой модели. "
"Сожми старую часть диалога в 58 коротких пунктов на русском. "
"Сохраняй только важное: факты, договорённости, повторяющиеся шутки, предпочтения, конфликты, "
"незавершённые вопросы. Не выдумывай. Ответь только итоговой сводкой."
)
USER_FALLBACK_TEXT = "Ты чё, кент? Напиши текст, а не пустоту."
EMPTY_RESPONSE_TEXT = "Братуха, чёт базар не клеится, попробуй ещё раз."
ERROR_RESPONSE_TEXT = "Бля, кент, чёт движок заглох. Попробуй позже."
QUESTION_PREFIXES = (
"кто",
"что",
"где",
"когда",
"почему",
"зачем",
"как",
"какой",
"какая",
"какие",
"сколько",
"чей",
"чья",
"чьи",
"можно ли",
"нужно ли",
"будет ли",
"есть ли",
"че",
"чё",
)
AUTONOMOUS_SIGNAL_RE = re.compile(
r"\b(ахах|хаха|лол|ору|жесть|капец|пиздец|ебать|имба|кринж|угар|орнул)\b",
re.IGNORECASE,
)
LETTER_RE = re.compile(r"[A-Za-zА-Яа-яЁё0-9]")
_reply_lock = asyncio.Lock()
_summary_lock = asyncio.Lock()
_cached_bot_id: int | None = None
_cached_bot_username: str = ""
def _connect_db() -> sqlite3.Connection:
conn = sqlite3.connect(DB_PATH, timeout=30)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA busy_timeout = 30000")
return conn
def _column_exists(conn: sqlite3.Connection, table: str, column: str) -> bool:
rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
return any(row["name"] == column for row in rows)
def _get_state_from_conn(conn: sqlite3.Connection, chat_id: int, key: str, default: str = "") -> str:
row = conn.execute(
"SELECT value FROM chat_state WHERE chat_id = ? AND key = ?",
(chat_id, key),
).fetchone()
return row["value"] if row else default
def _set_state_from_conn(conn: sqlite3.Connection, chat_id: int, key: str, value: str) -> None:
conn.execute(
"""
INSERT INTO chat_state (chat_id, key, value)
VALUES (?, ?, ?)
ON CONFLICT(chat_id, key) DO UPDATE SET value = excluded.value
""",
(chat_id, key, value),
)
def _init_db() -> None:
Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True)
with _connect_db() as conn:
conn.execute("PRAGMA journal_mode = WAL")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
chat_id INTEGER NOT NULL,
role TEXT NOT NULL,
name TEXT NOT NULL,
text TEXT NOT NULL,
created_at REAL NOT NULL DEFAULT 0
)
"""
)
if not _column_exists(conn, "chat_history", "created_at"):
conn.execute("ALTER TABLE chat_history ADD COLUMN created_at REAL NOT NULL DEFAULT 0")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_state (
chat_id INTEGER NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
PRIMARY KEY (chat_id, key)
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_history_chat_id_id ON chat_history(chat_id, id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_history_chat_id_role_id ON chat_history(chat_id, role, id)"
)
conn.commit()
def _clean_text(text: str) -> str:
return (text or "").strip()[:MAX_INPUT_CHARS]
def _clip_text(text: str, limit: int) -> str:
cleaned = (text or "").strip()
if len(cleaned) <= limit:
return cleaned
return f"{cleaned[: max(0, limit - 1)].rstrip()}"
def push_message(chat_id: int, role: str, name: str, text: str) -> None:
cleaned_text = _clean_text(text)
if not cleaned_text:
return
with _connect_db() as conn:
conn.execute(
"""
INSERT INTO chat_history (chat_id, role, name, text, created_at)
VALUES (?, ?, ?, ?, ?)
""",
(chat_id, role, name, cleaned_text, time.time()),
)
conn.commit()
def _get_summary(chat_id: int) -> str:
with _connect_db() as conn:
return _get_state_from_conn(conn, chat_id, "summary", "")
def _get_history_rows(chat_id: int) -> list[sqlite3.Row]:
with _connect_db() as conn:
return conn.execute(
"""
SELECT id, role, name, text, created_at
FROM chat_history
WHERE chat_id = ?
ORDER BY id ASC
""",
(chat_id,),
).fetchall()
def _latest_reply_stats(chat_id: int) -> tuple[float, int]:
with _connect_db() as conn:
last_assistant = conn.execute(
"""
SELECT id, created_at
FROM chat_history
WHERE chat_id = ? AND role = 'assistant'
ORDER BY id DESC
LIMIT 1
""",
(chat_id,),
).fetchone()
if not last_assistant:
return 0.0, 10_000
user_messages_since_reply = conn.execute(
"""
SELECT COUNT(*)
FROM chat_history
WHERE chat_id = ? AND role = 'user' AND id > ?
""",
(chat_id, last_assistant["id"]),
).fetchone()[0]
return float(last_assistant["created_at"] or 0.0), int(user_messages_since_reply)
def _store_summary_and_prune(chat_id: int, summary: str, last_row_id: int) -> None:
with _connect_db() as conn:
_set_state_from_conn(conn, chat_id, "summary", summary)
conn.execute(
"DELETE FROM chat_history WHERE chat_id = ? AND id <= ?",
(chat_id, last_row_id),
)
conn.commit()
def _format_row_for_llm(row: sqlite3.Row) -> dict[str, str]:
if row["role"] == "user":
return {
"role": "user",
"content": f"{row['name']}: {_clip_text(row['text'], HISTORY_LINE_CHAR_LIMIT)}",
}
return {"role": "assistant", "content": _clip_text(row["text"], HISTORY_LINE_CHAR_LIMIT)}
def _build_messages(
chat_id: int,
system_prompt: str,
current_text: str | None = None,
current_name: str | None = None,
) -> list[dict[str, str]]:
current_payload = None
current_budget = 0
if current_text:
current_payload = {
"role": "user",
"content": f"{current_name or 'кент'}: {_clean_text(current_text)}",
}
current_budget = len(current_payload["content"])
summary = _get_summary(chat_id).strip()
summary_block = ""
if summary:
summary_block = f"Краткая память чата:\n{summary[:SUMMARY_CHAR_BUDGET]}"
used_chars = len(system_prompt) + len(summary_block) + current_budget
recent_messages: list[dict[str, str]] = []
for row in reversed(_get_history_rows(chat_id)[-RECENT_MESSAGES_LIMIT:]):
llm_message = _format_row_for_llm(row)
if used_chars + len(llm_message["content"]) > PROMPT_CHAR_BUDGET:
break
recent_messages.append(llm_message)
used_chars += len(llm_message["content"])
messages = [{"role": "system", "content": system_prompt}]
if summary_block:
messages.append({"role": "system", "content": summary_block})
messages.extend(reversed(recent_messages))
if current_payload:
messages.append(current_payload)
return messages
async def _call_llm(
messages: list[dict[str, str]],
*,
max_tokens: int,
temperature: float,
top_p: float,
) -> str:
url = f"{LLAMA_API_URL.rstrip('/')}/v1/chat/completions"
payload = {
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
"top_p": top_p,
}
timeout = aiohttp.ClientTimeout(total=120)
async with aiohttp.ClientSession(timeout=timeout) as session:
async with session.post(url, json=payload) as resp:
raw_text = await resp.text()
if resp.status >= 400:
logger.error("LLM API returned status %s: %s", resp.status, _clip_text(raw_text, 300))
raise RuntimeError(f"LLM API error {resp.status}: {raw_text[:300]}")
try:
data = await resp.json(content_type=None)
except Exception as exc:
logger.error("LLM API returned invalid JSON: %s", _clip_text(raw_text, 300))
raise RuntimeError(f"Invalid LLM API response: {raw_text[:300]}") from exc
choices = data.get("choices") or []
if not choices:
logger.error("LLM API returned no choices: %s", _clip_text(str(data), 300))
raise RuntimeError(f"LLM API returned no choices: {data}")
content = choices[0].get("message", {}).get("content", "")
cleaned_content = content.strip()
if not cleaned_content:
logger.warning("LLM returned empty content. Raw response: %s", _clip_text(raw_text, 300))
return cleaned_content
async def _maybe_refresh_summary(chat_id: int) -> None:
async with _summary_lock:
# Сворачиваем старую часть истории в summary, чтобы не пихать весь лог в модель.
for _ in range(6):
rows = await asyncio.to_thread(_get_history_rows, chat_id)
if len(rows) <= SUMMARY_TRIGGER_MESSAGES:
return
available_to_summarize = len(rows) - RECENT_MESSAGES_LIMIT
if available_to_summarize <= 0:
return
batch_size = min(available_to_summarize, SUMMARY_BATCH_MESSAGES)
rows_to_summarize = rows[:batch_size]
current_summary = await asyncio.to_thread(_get_summary, chat_id)
transcript = "\n".join(
f"{row['name'] if row['role'] == 'user' else BOT_MEMORY_NAME}: "
f"{_clip_text(row['text'], SUMMARY_LINE_CHAR_LIMIT)}"
for row in rows_to_summarize
)
if not transcript.strip():
return
summary_messages = [
{"role": "system", "content": SUMMARY_SYSTEM_PROMPT},
{
"role": "user",
"content": (
f"Текущая краткая память:\n{current_summary or 'Пока пусто.'}\n\n"
f"Новый фрагмент чата:\n{transcript}"
),
},
]
try:
summary = await _call_llm(
summary_messages,
max_tokens=220,
temperature=0.2,
top_p=0.9,
)
except Exception:
logger.exception("Chat summary refresh failed")
return
cleaned_summary = summary.strip()
if not cleaned_summary:
logger.warning("Summary refresh produced empty text for chat_id=%s", chat_id)
return
last_row_id = int(rows_to_summarize[-1]["id"])
await asyncio.to_thread(_store_summary_and_prune, chat_id, cleaned_summary, last_row_id)
async def _generate_response(
chat_id: int,
*,
system_prompt: str,
current_text: str | None = None,
current_name: str | None = None,
max_tokens: int,
temperature: float,
top_p: float,
) -> str:
await _maybe_refresh_summary(chat_id)
messages = await asyncio.to_thread(
_build_messages,
chat_id,
system_prompt,
current_text,
current_name,
)
return await _call_llm(
messages,
max_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
)
async def _get_bot_identity(message: Message) -> tuple[int | None, str]:
global _cached_bot_id, _cached_bot_username
if _cached_bot_id is None:
me = await message.bot.get_me()
_cached_bot_id = me.id
_cached_bot_username = (me.username or "").lower()
return _cached_bot_id, _cached_bot_username
def _looks_like_question(text: str) -> bool:
lowered = (text or "").strip().lower()
if not lowered:
return False
if "?" in lowered or "" in lowered:
return True
normalized = re.sub(r"^[^a-zа-яё0-9]+", "", lowered)
return any(normalized == prefix or normalized.startswith(f"{prefix} ") for prefix in QUESTION_PREFIXES)
def _is_reply_to_bot(message: Message, bot_id: int | None) -> bool:
reply = message.reply_to_message
return bool(bot_id and reply and reply.from_user and reply.from_user.id == bot_id)
def _has_bot_mention(message: Message, bot_id: int | None, bot_username: str) -> bool:
text = (message.text or "").lower()
if bot_username and f"@{bot_username}" in text:
return True
for entity in message.entities or []:
entity_type = getattr(entity.type, "value", entity.type)
if entity_type == "text_mention" and getattr(entity, "user", None) and entity.user.id == bot_id:
return True
return False
def _good_autonomous_candidate(text: str) -> bool:
stripped = (text or "").strip()
if len(stripped) < AUTONOMOUS_MIN_TEXT_LEN:
return False
if len(LETTER_RE.findall(stripped)) < 10:
return False
if AUTONOMOUS_SIGNAL_RE.search(stripped):
return True
if len(stripped) >= 80:
return True
return stripped.count("!") >= 2 or "..." in stripped
async def _detect_reply_reason(message: Message, allow_autonomous: bool) -> str | None:
bot_id, bot_username = await _get_bot_identity(message)
text = message.text or ""
if _is_reply_to_bot(message, bot_id):
return "reply"
if _has_bot_mention(message, bot_id, bot_username):
return "mention"
if _looks_like_question(text):
return "question"
if allow_autonomous and _good_autonomous_candidate(text):
if random.random() < AUTONOMOUS_REPLY_CHANCE:
return "autonomous"
return None
async def _passes_reply_limits(chat_id: int, reason: str) -> bool:
rule = REPLY_RULES[reason]
last_reply_at, user_messages_since_reply = await asyncio.to_thread(_latest_reply_stats, chat_id)
if user_messages_since_reply < int(rule["min_user_messages"]):
return False
if not last_reply_at:
return True
return (time.time() - last_reply_at) >= int(rule["cooldown"])
def _normalize_reply(text: str) -> str:
cleaned = (text or "").strip()
if not cleaned:
return ""
if cleaned.lower().startswith(SKIP_TOKEN.lower()):
return ""
return cleaned
async def handle_chat_message(message: Message, *, store_message: bool = True, allow_autonomous: bool = True) -> bool:
if not message.text or message.text.startswith("/"):
return False
if not message.from_user or message.from_user.is_bot:
return False
chat_id = message.chat.id
user_name = message.from_user.first_name or "кент"
user_text = _clean_text(message.text)
if not user_text:
return False
if store_message:
await asyncio.to_thread(push_message, chat_id, "user", user_name, user_text)
# Один чат, поэтому ответы сериализуем и не даём боту наспамить параллельными реплаями.
async with _reply_lock:
reason = await _detect_reply_reason(message, allow_autonomous=allow_autonomous)
if not reason or not await _passes_reply_limits(chat_id, reason):
return False
system_prompt = AUTOREPLY_SYSTEM_PROMPT if reason == "autonomous" else SYSTEM_PROMPT
max_tokens = 120 if reason == "autonomous" else 220
temperature = 0.9 if reason == "autonomous" else 0.8
try:
await message.bot.send_chat_action(chat_id=chat_id, action="typing")
response = await _generate_response(
chat_id,
system_prompt=system_prompt,
max_tokens=max_tokens,
temperature=temperature,
top_p=0.9,
)
except Exception:
logger.exception("Chat reply generation failed")
return False
normalized_response = _normalize_reply(response)
if not normalized_response:
logger.warning(
"Reply skipped after normalization. chat_id=%s reason=%s raw=%s",
chat_id,
reason,
_clip_text(response, 200),
)
return False
await asyncio.to_thread(push_message, chat_id, "assistant", BOT_MEMORY_NAME, normalized_response)
await message.reply(normalized_response, parse_mode=None)
return True
async def generate_autoreply(chat_id: int, text: str, user_name: str) -> str:
response = await _generate_response(
chat_id,
system_prompt=AUTOREPLY_SYSTEM_PROMPT,
current_text=text,
current_name=user_name,
max_tokens=120,
temperature=0.9,
top_p=0.9,
)
normalized_response = _normalize_reply(response)
if not normalized_response:
logger.warning(
"Legacy generate_autoreply produced empty/skip response. chat_id=%s raw=%s",
chat_id,
_clip_text(response, 200),
)
return normalized_response
async def handle_talk(message: Message) -> None:
parts = (message.text or "").split(maxsplit=1)
if len(parts) < 2 or not parts[1].strip():
await message.reply("Ты чё, кент? Напиши /talk [текст], побазарим.")
return
if not message.from_user:
await message.reply(USER_FALLBACK_TEXT)
return
chat_id = message.chat.id
user_name = message.from_user.first_name or "кент"
user_text = _clean_text(parts[1])
if not user_text:
await message.reply(USER_FALLBACK_TEXT)
return
await asyncio.to_thread(push_message, chat_id, "user", user_name, user_text)
async with _reply_lock:
try:
await message.bot.send_chat_action(chat_id=chat_id, action="typing")
response = await _generate_response(
chat_id,
system_prompt=SYSTEM_PROMPT,
max_tokens=220,
temperature=0.8,
top_p=0.9,
)
except Exception:
logger.exception("Talk command generation failed")
await message.reply(ERROR_RESPONSE_TEXT)
return
normalized_response = _normalize_reply(response)
if not normalized_response:
logger.warning(
"Talk command produced empty/skip response. chat_id=%s user=%s raw=%s",
chat_id,
user_name,
_clip_text(response, 200),
)
normalized_response = EMPTY_RESPONSE_TEXT
await asyncio.to_thread(push_message, chat_id, "assistant", BOT_MEMORY_NAME, normalized_response)
await message.reply(normalized_response, parse_mode=None)
_init_db()