Add chat and image message handlers

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}")