Add chat and image message handlers
This commit is contained in:
arsam
@@ -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}")
|
||||
Reference in New Issue
Block a user