From d2f106cc9e81509cbbc697982230324225ebfb4f Mon Sep 17 00:00:00 2001 From: arsam Date: Fri, 18 Sep 2026 22:41:18 +0330 Subject: [PATCH] Add chat and image message handlers --- handlers/chat.py | 290 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 290 insertions(+) create mode 100644 handlers/chat.py diff --git a/handlers/chat.py b/handlers/chat.py new file mode 100644 index 0000000..290d2b4 --- /dev/null +++ b/handlers/chat.py @@ -0,0 +1,290 @@ +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}")