import asyncio import re import time from abc import ABC from typing import Union from grpc import StatusCode from grpc.experimental.aio import AioRpcError from library.telegram.base import RequestContext from library.telegram.common import close_button from library.telegram.utils import safe_execution from nexus.bot.exceptions import BannedUserError from nexus.bot.widgets.search_widget import ( InlineSearchWidget, SearchWidget, ) from nexus.translations import t from nexus.views.telegram.base_holder import BaseHolder from nexus.views.telegram.common import encode_deep_query from telethon import ( Button, events, ) from telethon.tl.types import InlineQueryPeerTypeSameBotPM from .base import ( BaseCallbackQueryHandler, BaseHandler, ) class BaseSearchHandler(BaseHandler, ABC): async def setup_widget( self, request_context: RequestContext, prefetch_message, query: str, is_shortpath_enabled: bool = False, ) -> tuple[str, list[Union[list[Button]], list[Button]]]: session_id = self.generate_session_id() message_id = prefetch_message.id request_context.add_default_fields( is_group_mode=request_context.is_group_mode(), mode='search', session_id=session_id, ) start_time = time.time() language = request_context.chat.language bot_name = self.application.config['telegram']['bot_name'] try: search_widget = await SearchWidget.create( application=self.application, chat=request_context.chat, session_id=session_id, request_id=request_context.request_id, query=query, is_group_mode=request_context.is_group_mode(), ) except AioRpcError as e: if e.code() == StatusCode.INVALID_ARGUMENT: return t('INVALID_SYNTAX_ERROR', language).format( too_difficult_picture_url=self.application.config['application'].get('too_difficult_picture_url', ''), ), [close_button()] elif e.code() == StatusCode.CANCELLED: return t('MAINTENANCE', language).format( maintenance_picture_url=self.application.config['application'].get('maintenance_picture_url', ''), ), [close_button()], request_context.error_log(e) raise e request_context.statbox( action='documents_retrieved', duration=time.time() - start_time, query=query, page=0, scored_documents=len(search_widget.scored_documents), ) if len(search_widget.scored_documents) == 1 and is_shortpath_enabled: holder = BaseHolder.create(search_widget.scored_documents[0].typed_document) view = holder.view_builder(language).add_view(bot_name=bot_name).build() buttons = holder.buttons_builder(language).add_default_layout( bot_name=bot_name, session_id=session_id, position=0, ).build() return view, buttons return await search_widget.render(message_id=message_id) class SearchHandler(BaseSearchHandler): filter = events.NewMessage(incoming=True, pattern=re.compile(r'^(/search(?:@\w+)?\s+)?(.*)', flags=re.DOTALL)) is_group_handler = True should_reset_last_widget = False is_subscription_required_for_handler = True def check_search_ban_timeout(self, user_id: str): ban_timeout = self.application.user_manager.check_search_ban_timeout(user_id=user_id) if ban_timeout: raise BannedUserError(ban_timeout=ban_timeout) self.application.user_manager.add_search_time(user_id=user_id, search_time=time.time()) def parse_pattern(self, event: events.ChatAction): search_prefix = event.pattern_match.group(1) query = event.pattern_match.group(2).strip() return search_prefix, query async def handler(self, event: events.ChatAction, request_context: RequestContext): language = request_context.chat.language try: self.check_search_ban_timeout(user_id=str(request_context.chat.chat_id)) except BannedUserError as e: request_context.error_log(e) async with safe_execution(error_log=request_context.error_log): return await event.reply(t('BANNED_FOR_SECONDS', language).format( seconds=e.ban_timeout, reason=t('BAN_MESSAGE_TOO_MANY_REQUESTS', language), )) search_prefix, query = self.parse_pattern(event) if request_context.is_group_mode() and not search_prefix: return if request_context.is_personal_mode() and search_prefix: query = event.raw_text prefetch_message = await event.reply( t("SEARCHING", language), ) self.application.user_manager.last_widget[request_context.chat.chat_id] = prefetch_message.id try: text, buttons = await self.setup_widget( request_context=request_context, prefetch_message=prefetch_message, query=query, is_shortpath_enabled=True, ) return await self.application.telegram_client.edit_message( request_context.chat.chat_id, prefetch_message.id, text, buttons=buttons, link_preview=False, ) except (AioRpcError, asyncio.CancelledError) as e: await asyncio.gather( event.delete(), prefetch_message.delete(), ) raise e class InlineSearchHandler(BaseSearchHandler): filter = events.InlineQuery() stop_propagation = False async def handler(self, event, request_context: RequestContext): if event.query.peer_type == InlineQueryPeerTypeSameBotPM(): await event.answer() return builder = event.builder session_id = self.generate_session_id() try: if len(event.text) <= 3: await event.answer([]) raise events.StopPropagation() inline_search_widget = await InlineSearchWidget.create( application=self.application, chat=request_context.chat, session_id=session_id, request_id=request_context.request_id, query=event.text, is_group_mode=request_context.is_group_mode(), ) items = inline_search_widget.render(builder=builder) encoded_query = encode_deep_query(event.text) if len(encoded_query) < 32: await event.answer( items, private=True, switch_pm=self.application.config['telegram']['bot_name'], switch_pm_param=encoded_query, ) else: await event.answer(items) except AioRpcError as e: if e.code() == StatusCode.INVALID_ARGUMENT or e.code() == StatusCode.CANCELLED: await event.answer([]) raise e raise events.StopPropagation() class SearchEditHandler(BaseSearchHandler): filter = events.MessageEdited(incoming=True, pattern=re.compile(r'^(/search(?:@\w+)\s+)?(.*)', flags=re.DOTALL)) is_group_handler = True should_reset_last_widget = False def parse_pattern(self, event: events.ChatAction): search_prefix = event.pattern_match.group(1) query = event.pattern_match.group(2).strip() return search_prefix, query async def handler(self, event: events.ChatAction, request_context: RequestContext): search_prefix, query = self.parse_pattern(event) request_context.add_default_fields(mode='search_edit') if request_context.is_group_mode() and not search_prefix: return if request_context.is_personal_mode() and search_prefix: query = event.raw_text for next_message in await self.get_last_messages_in_chat(event): if next_message.is_reply and event.id == next_message.reply_to_msg_id: request_context.statbox(action='resolved') text, buttons = await self.setup_widget( request_context=request_context, prefetch_message=next_message, query=query, ) return await self.application.telegram_client.edit_message( request_context.chat.chat_id, next_message.id, text, buttons=buttons, link_preview=False, ) return await event.reply( t('REPLY_MESSAGE_HAS_BEEN_DELETED', request_context.chat.language), ) class SearchPagingHandler(BaseCallbackQueryHandler): filter = events.CallbackQuery(pattern='^/search_([A-Za-z0-9]+)_([0-9]+)_([0-9]+)$') should_reset_last_widget = False def parse_pattern(self, event: events.ChatAction): session_id = event.pattern_match.group(1).decode() message_id = int(event.pattern_match.group(2).decode()) page = int(event.pattern_match.group(3).decode()) return session_id, message_id, page async def handler(self, event: events.ChatAction, request_context: RequestContext): session_id, message_id, page = self.parse_pattern(event) request_context.add_default_fields(mode='search_paging', session_id=session_id) start_time = time.time() message = await event.get_message() if not message: return await event.answer() reply_message = await message.get_reply_message() if not reply_message: return await event.respond( t('REPLY_MESSAGE_HAS_BEEN_DELETED', request_context.chat.language), ) query = reply_message.raw_text.replace(f'@{self.application.config["telegram"]["bot_name"]}', '').strip() try: search_widget = await SearchWidget.create( application=self.application, chat=request_context.chat, session_id=session_id, request_id=request_context.request_id, query=query, page=page, ) except AioRpcError as e: if e.code() == StatusCode.INVALID_ARGUMENT or e.code() == StatusCode.CANCELLED: request_context.error_log(e) return await event.answer( t('MAINTENANCE_WO_PIC', request_context.chat.language), ) raise e request_context.statbox( action='documents_retrieved', duration=time.time() - start_time, query=query, page=page, scored_documents=len(search_widget.scored_documents), ) serp, buttons = await search_widget.render(message_id=message_id) return await asyncio.gather( event.answer(), message.edit(serp, buttons=buttons, link_preview=False) )