291 lines
10 KiB
Python
291 lines
10 KiB
Python
import re
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from telegram import Update, Message
|
|
from telegram.constants import ChatAction, ChatType
|
|
from telegram.ext import ContextTypes
|
|
from services import auth, ai_client
|
|
from utils import context, chunker, persian, formatter
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_MEDIA_GROUP_WAIT = 1.2
|
|
|
|
|
|
# ── Sending helpers ───────────────────────────────────────────
|
|
|
|
async def _send(message: Message, text: str):
|
|
"""Send text with HTML formatting; fall back to plain on parse error."""
|
|
formatted, mode = formatter.safe_send(text)
|
|
for chunk in chunker.chunk_text(formatted):
|
|
try:
|
|
await message.reply_text(chunk, parse_mode=mode)
|
|
except Exception:
|
|
plain = re.sub(r"<[^>]+>", "", chunk)
|
|
await message.reply_text(plain)
|
|
|
|
|
|
async def _send_images(message: Message, urls: list[str]):
|
|
"""Upload up to 5 image URLs to Telegram."""
|
|
sent = 0
|
|
for url in urls[:5]:
|
|
try:
|
|
await message.reply_photo(url)
|
|
sent += 1
|
|
except Exception as exc:
|
|
logger.warning("Could not send image %s: %s", url[:60], exc)
|
|
return sent
|
|
|
|
|
|
# ── Group-mention / reply detection ──────────────────────────
|
|
|
|
def _is_addressed(message: Message, bot_username: str, bot_id: int) -> bool:
|
|
text = message.text or message.caption or ""
|
|
|
|
# Reply to bot's own message
|
|
rt = message.reply_to_message
|
|
if rt and rt.from_user:
|
|
if rt.from_user.id == bot_id:
|
|
return True
|
|
if rt.from_user.username and rt.from_user.username.lower() == bot_username.lower():
|
|
return True
|
|
|
|
# @mention in entities
|
|
for entity in list(message.entities or []) + list(message.caption_entities or []):
|
|
if entity.type == "mention":
|
|
mention = text[entity.offset: entity.offset + entity.length]
|
|
if mention.lower() == f"@{bot_username.lower()}":
|
|
return True
|
|
|
|
# Name mention (Puppy / پاپی)
|
|
return persian.is_puppy_mentioned(text)
|
|
|
|
|
|
def _clean_text(text: str, bot_username: str) -> str:
|
|
text = re.sub(rf"@{re.escape(bot_username)}\s*", "", text, flags=re.IGNORECASE)
|
|
text = persian.strip_puppy_prefix(text)
|
|
return text.strip()
|
|
|
|
|
|
def _get_quoted_context(message: Message, bot_id: int) -> str:
|
|
"""If user replied to another user's message, return a quoted prefix for context."""
|
|
rt = message.reply_to_message
|
|
if not rt or not rt.from_user:
|
|
return ""
|
|
if rt.from_user.id == bot_id:
|
|
return "" # replying to bot is handled naturally via history
|
|
quoted = rt.text or rt.caption or ""
|
|
if not quoted:
|
|
return ""
|
|
name = rt.from_user.first_name or "someone"
|
|
return f'[Replying to {name}: "{quoted[:300]}"]\n'
|
|
|
|
|
|
def _image_analysis_mode(caption: str) -> str:
|
|
text = caption.lower()
|
|
if any(word in text for word in ("ocr", "text", "read", "extract", "متن", "بخوان")):
|
|
return "ocr"
|
|
if any(word in text for word in ("describe", "description", "توصیف", "شرح")):
|
|
return "describe"
|
|
if any(word in text for word in ("objects", "things", "items", "اشیا", "چیزها")):
|
|
return "objects"
|
|
if "?" in caption or any(word in text for word in ("what", "why", "how", "where", "who", "چی", "چرا", "کجا")):
|
|
return "question"
|
|
return "auto"
|
|
|
|
|
|
async def _download_image(message: Message, ctx: ContextTypes.DEFAULT_TYPE) -> tuple[bytes, str] | None:
|
|
if message.photo:
|
|
image_file = await ctx.bot.get_file(message.photo[-1].file_id)
|
|
return bytes(await image_file.download_as_bytearray()), "image/jpeg"
|
|
if message.document and (message.document.mime_type or "").startswith("image/"):
|
|
image_file = await ctx.bot.get_file(message.document.file_id)
|
|
return bytes(await image_file.download_as_bytearray()), message.document.mime_type or "image/jpeg"
|
|
return None
|
|
|
|
|
|
def _is_image_message(message: Message) -> bool:
|
|
return bool(message.photo or (message.document and (message.document.mime_type or "").startswith("image/")))
|
|
|
|
|
|
# ── Image handler ─────────────────────────────────────────────
|
|
|
|
async def _handle_photo(update: Update, ctx: ContextTypes.DEFAULT_TYPE, force_mode: str | None = None):
|
|
message = update.effective_message
|
|
user = update.effective_user
|
|
chat = update.effective_chat
|
|
|
|
caption = _clean_text(message.caption or "", ctx.bot.username or "")
|
|
settings = context.get_user_settings(user.id)
|
|
history = context.get_history(message.chat_id)
|
|
sys_prompt = context.get_system_prompt()
|
|
if context.reset_chat_if_input_too_large(message.chat_id, sys_prompt, keep_last=False):
|
|
history = []
|
|
|
|
await ctx.bot.send_chat_action(chat.id, ChatAction.TYPING)
|
|
|
|
try:
|
|
downloaded = await _download_image(message, ctx)
|
|
if not downloaded:
|
|
return
|
|
|
|
image_bytes, mime_type = downloaded
|
|
analysis_mode = force_mode or _image_analysis_mode(caption)
|
|
|
|
reply = await ai_client.chat_with_image(
|
|
messages=history,
|
|
system_prompt=sys_prompt,
|
|
max_tokens=settings["max_tokens"],
|
|
image_bytes=image_bytes,
|
|
mime_type=mime_type,
|
|
caption=caption,
|
|
analysis_mode=analysis_mode,
|
|
)
|
|
|
|
context.add_message(message.chat_id, "user", f"[Photo]{': ' + caption if caption else ''}")
|
|
context.add_message(message.chat_id, "assistant", reply)
|
|
context.increment_msg()
|
|
context.record_event("images_analyzed")
|
|
|
|
await _send(message, reply)
|
|
|
|
except Exception as exc:
|
|
context.record_error(f"Image error: {exc}")
|
|
await message.reply_text(f"❌ Image error: {exc}")
|
|
|
|
|
|
async def _handle_media_group(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
|
|
message = update.effective_message
|
|
group_id = message.media_group_id
|
|
if not group_id:
|
|
await _handle_photo(update, ctx)
|
|
return
|
|
|
|
groups = ctx.application.bot_data.setdefault("media_groups", {})
|
|
bucket = groups.setdefault(group_id, {"messages": [], "last_seen": 0.0})
|
|
bucket["messages"].append(message)
|
|
bucket["last_seen"] = time.time()
|
|
|
|
if len(bucket["messages"]) > 1:
|
|
return
|
|
|
|
await asyncio.sleep(_MEDIA_GROUP_WAIT)
|
|
bucket = groups.pop(group_id, bucket)
|
|
messages = bucket["messages"][:10]
|
|
first = messages[0]
|
|
user = first.from_user
|
|
settings = context.get_user_settings(user.id)
|
|
history = context.get_history(first.chat_id)
|
|
sys_prompt = context.get_system_prompt()
|
|
if context.reset_chat_if_input_too_large(first.chat_id, sys_prompt, keep_last=False):
|
|
history = []
|
|
caption = ""
|
|
images: list[tuple[bytes, str]] = []
|
|
|
|
await ctx.bot.send_chat_action(first.chat_id, ChatAction.TYPING)
|
|
|
|
try:
|
|
for msg in messages:
|
|
if not caption and msg.caption:
|
|
caption = _clean_text(msg.caption, ctx.bot.username or "")
|
|
downloaded = await _download_image(msg, ctx)
|
|
if downloaded:
|
|
images.append(downloaded)
|
|
|
|
if not images:
|
|
return
|
|
|
|
reply = await ai_client.chat_with_images(
|
|
messages=history,
|
|
system_prompt=sys_prompt,
|
|
max_tokens=settings["max_tokens"],
|
|
images=images,
|
|
caption=caption,
|
|
analysis_mode=_image_analysis_mode(caption),
|
|
)
|
|
|
|
context.add_message(first.chat_id, "user", f"[{len(images)} photos]{': ' + caption if caption else ''}")
|
|
context.add_message(first.chat_id, "assistant", reply)
|
|
context.increment_msg()
|
|
context.record_event("multi_image_analyses")
|
|
|
|
await _send(first, reply)
|
|
|
|
except Exception as exc:
|
|
context.record_error(f"Multi-image error: {exc}")
|
|
await first.reply_text(f"❌ Multi-image error: {exc}")
|
|
|
|
|
|
# ── Main message handler ──────────────────────────────────────
|
|
|
|
async def message_handler(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
|
|
message = update.effective_message
|
|
if not message:
|
|
return
|
|
|
|
user = update.effective_user
|
|
chat = update.effective_chat
|
|
|
|
if not auth.is_authorized(user.id, user.username or ""):
|
|
return
|
|
auth.register_known(user.id, username=user.username or "")
|
|
|
|
# Group gate: must be addressed
|
|
if chat.type in (ChatType.GROUP, ChatType.SUPERGROUP):
|
|
if not _is_addressed(message, ctx.bot.username or "", ctx.bot.id):
|
|
return
|
|
|
|
# Image branch: Telegram photos and image files/documents.
|
|
if _is_image_message(message):
|
|
if message.media_group_id:
|
|
await _handle_media_group(update, ctx)
|
|
else:
|
|
await _handle_photo(update, ctx)
|
|
return
|
|
|
|
text = message.text or ""
|
|
if not text:
|
|
return
|
|
|
|
text = _clean_text(text, ctx.bot.username or "")
|
|
if not text:
|
|
return
|
|
|
|
# Include quoted message from another user as context prefix
|
|
quoted = _get_quoted_context(message, ctx.bot.id)
|
|
full_text = quoted + text
|
|
|
|
settings = context.get_user_settings(user.id)
|
|
sys_prompt = context.get_system_prompt()
|
|
|
|
# Add to history first so AI sees the current message
|
|
context.add_message(message.chat_id, "user", full_text)
|
|
context.reset_chat_if_input_too_large(message.chat_id, sys_prompt, keep_last=True)
|
|
history = context.get_history(message.chat_id)
|
|
|
|
await ctx.bot.send_chat_action(chat.id, ChatAction.TYPING)
|
|
|
|
try:
|
|
reply, image_urls = await ai_client.chat(
|
|
messages=history,
|
|
system_prompt=sys_prompt,
|
|
max_tokens=settings["max_tokens"],
|
|
search_enabled=settings["search_enabled"],
|
|
)
|
|
|
|
context.add_message(message.chat_id, "assistant", reply)
|
|
context.increment_msg()
|
|
context.record_event("text_chats")
|
|
|
|
# Send images first (appear above text)
|
|
if image_urls:
|
|
await _send_images(message, image_urls)
|
|
|
|
await _send(message, reply)
|
|
|
|
except Exception as exc:
|
|
context.pop_last_message(message.chat_id)
|
|
context.record_error(f"Chat error: {exc}")
|
|
logger.error("Chat error: %s", exc)
|
|
await message.reply_text(f"❌ Error: {exc}")
|