hyperboria/nexus/bot/handlers/search.py
the-superpirate 681817ceae No description
GitOrigin-RevId: 83514338be1d662518bab9fe5ab6eefef98f229f
2021-04-09 13:29:10 +03:00

290 lines
11 KiB
Python

import asyncio
import logging
import re
import time
from grpc import StatusCode
from grpc.experimental.aio import AioRpcError
from library.telegram.base import RequestContext
from nexus.bot.widgets.search_widget import SearchWidget
from nexus.translations import t
from nexus.views.telegram.common import close_button
from nexus.views.telegram.registry import parse_typed_document_to_view
from telethon import (
events,
functions,
)
from .base import (
BaseCallbackQueryHandler,
BaseHandler,
)
class BaseSearchHandler(BaseHandler):
should_reset_last_widget = False
async def do_search(
self,
event: events.ChatAction,
request_context: RequestContext,
prefetch_message,
query: str,
is_group_mode: bool = False,
is_shortpath_enabled: bool = False,
):
session_id = self.generate_session_id()
message_id = prefetch_message.id
request_context.add_default_fields(is_group_mode=is_group_mode, mode='search', session_id=session_id)
start_time = time.time()
try:
search_widget = await SearchWidget.create(
application=self.application,
chat=request_context.chat,
session_id=session_id,
message_id=message_id,
request_id=request_context.request_id,
query=query,
is_group_mode=is_group_mode,
)
except AioRpcError as e:
actions = [
self.application.telegram_client.delete_messages(
request_context.chat.id,
[message_id],
)
]
if e.code() == StatusCode.INVALID_ARGUMENT:
too_difficult_picture_url = self.application.config['application'].get('too_difficult_picture_url', '')
if e.details() == 'url_query_error':
actions.append(
event.reply(
t('INVALID_QUERY_ERROR', language=request_context.chat.language).format(
too_difficult_picture_url=too_difficult_picture_url,
),
buttons=[close_button()],
)
)
elif e.details() == 'invalid_query_error':
actions.append(
event.reply(
t('INVALID_SYNTAX_ERROR', language=request_context.chat.language).format(
too_difficult_picture_url=too_difficult_picture_url,
),
buttons=[close_button()],
)
)
return await asyncio.gather(*actions)
elif e.code() == StatusCode.CANCELLED:
maintenance_picture_url = self.application.config['application'].get('maintenance_picture_url', '')
request_context.error_log(e)
actions.append(event.reply(
t('MAINTENANCE', language=request_context.chat.language).format(
maintenance_picture_url=maintenance_picture_url,
),
buttons=[close_button()],
))
return await asyncio.gather(*actions)
raise e
action = 'documents_found'
if len(search_widget.scored_documents) == 0:
action = 'documents_not_found'
request_context.statbox(
action=action,
duration=time.time() - start_time,
query=f'page:0 query:{query}',
)
if len(search_widget.scored_documents) == 1 and is_shortpath_enabled:
scored_document = search_widget.scored_documents[0]
document_view = parse_typed_document_to_view(scored_document.typed_document)
# Second (re-)fetching is required to retrieve duplicates
document_view = await self.resolve_document(
schema=scored_document.typed_document.WhichOneof('document'),
document_id=document_view.id,
position=0,
session_id=session_id,
request_context=request_context,
)
view, buttons = document_view.get_view(
language=request_context.chat.language,
session_id=session_id,
bot_external_name=self.application.config['telegram']['bot_external_name'],
with_buttons=not is_group_mode,
)
return await asyncio.gather(
self.application.telegram_client.edit_message(
request_context.chat.id,
message_id,
view,
buttons=buttons,
),
)
serp, buttons = await search_widget.render()
return await self.application.telegram_client.edit_message(
request_context.chat.id,
message_id,
serp,
buttons=buttons,
link_preview=False,
)
class SearchHandler(BaseSearchHandler):
filter = events.NewMessage(incoming=True, pattern=re.compile('^(/search\\s+)?(.*)', flags=re.DOTALL))
is_group_handler = True
should_reset_last_widget = False
is_subscription_required_for_handler = True
async def ban_handler(self, event: events.ChatAction, request_context: RequestContext, ban_timeout: float):
logging.getLogger('statbox').info({
'bot_name': self.application.config['telegram']['bot_name'],
'action': 'user_flood_ban',
'mode': 'search',
'ban_timeout_seconds': ban_timeout,
'chat_id': request_context.chat.id,
})
ban_reason = t(
'BAN_MESSAGE_TOO_MANY_REQUESTS',
language=request_context.chat.language
)
return await event.reply(t(
'BANNED_FOR_SECONDS',
language=request_context.chat.language
).format(
seconds=str(ban_timeout),
reason=ban_reason,
))
async def handler(self, event: events.ChatAction, request_context: RequestContext):
ban_timeout = self.application.user_manager.check_search_ban_timeout(user_id=request_context.chat.id)
if ban_timeout:
return await self.ban_handler(event, request_context, ban_timeout)
self.application.user_manager.add_search_time(user_id=request_context.chat.id, search_time=time.time())
search_prefix = event.pattern_match.group(1)
query = event.pattern_match.group(2)
is_group_mode = event.is_group or event.is_channel
if is_group_mode and not search_prefix:
return
if not is_group_mode and search_prefix:
query = event.raw_text
prefetch_message = await event.reply(
t("SEARCHING", language=request_context.chat.language),
)
self.application.user_manager.last_widget[request_context.chat.id] = prefetch_message.id
try:
await self.do_search(
event, request_context, prefetch_message,
query=query,
is_group_mode=is_group_mode,
is_shortpath_enabled=True,
)
except (AioRpcError, asyncio.CancelledError) as e:
await asyncio.gather(
event.delete(),
prefetch_message.delete(),
)
raise e
class SearchEditHandler(BaseSearchHandler):
filter = events.MessageEdited(incoming=True, pattern=re.compile('^(/search\\s+)?(.*)', flags=re.DOTALL))
is_group_handler = True
should_reset_last_widget = False
async def handler(self, event: events.ChatAction, request_context: RequestContext):
request_context.add_default_fields(mode='search_edit')
search_prefix = event.pattern_match.group(1)
query = event.pattern_match.group(2)
is_group_mode = event.is_group or event.is_channel
if is_group_mode and not search_prefix:
return
if not is_group_mode and search_prefix:
query = event.raw_text
result = await self.application.telegram_client(functions.messages.GetMessagesRequest(
id=list(range(event.id + 1, event.id + 10)))
)
if not result:
request_context.statbox(action='failed')
return await event.reply(
t('REPLY_MESSAGE_HAS_BEEN_DELETED', language=request_context.chat.language),
)
for next_message in result.messages:
if next_message.is_reply and event.id == next_message.reply_to_msg_id:
request_context.statbox(action='resolved')
await self.do_search(
event,
request_context,
prefetch_message=next_message,
query=query,
is_group_mode=is_group_mode,
)
return
request_context.statbox(action='failed')
return await event.reply(
t('REPLY_MESSAGE_HAS_BEEN_DELETED', language=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
async def handler(self, event: events.ChatAction, request_context: RequestContext):
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())
request_context.add_default_fields(mode='search_paging', session_id=session_id)
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', language=request_context.chat.language),
)
start_time = time.time()
query = reply_message.raw_text
try:
search_widget = await SearchWidget.create(
application=self.application,
chat=request_context.chat,
session_id=session_id,
message_id=message_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', language=request_context.chat.language),
)
raise e
action = 'documents_found'
if len(search_widget.scored_documents) == 0:
action = 'documents_not_found'
request_context.statbox(
action=action,
duration=time.time() - start_time,
query=f'page:{page} query:{query}',
)
serp, buttons = await search_widget.render()
return await asyncio.gather(
event.answer(),
message.edit(serp, buttons=buttons, link_preview=False)
)