import time import logging import datetime import asyncio from sqlalchemy import select, func, and_ import interactions from interactions import ( Intents, Message, AutoArchiveDuration, ChannelType, MessageType, ActionRow, Button, ButtonStyle, ) from interactions.api.events import Component, MessageCreate from community_qa_bot.settings import ( DEBUG, TIDB_AI_API_KEY, TIDB_AI_DISCORD_CHAT_ENGINE, DISCORD_BOT_TOKEN, DISCORD_HELP_TEXT, DISCORD_HELP_CHANNEL_IDS, DISCORD_HELP_CHANNEL_NEW_QA_SERVICE_IDS, DISCORD_QA_DAILY_LIMIT, DISCORD_REACH_LIMIT_TEXT, DISCORD_NEW_HELP_MESSAGE_ALERT_CHANNEL_IDS, DISCORD_NEW_HELP_MESSAGE_FEISHU_WEBHOOK, ) # from community_qa_bot.pingcap_qa_service import QAService from community_qa_bot.pingcap_qa_service import CommunityQAService, CommunityNewQaService from community_qa_bot.db import AsyncSession as DBSession from community_qa_bot.models import ( QaHistory, User, QaFeedback, ) from community_qa_bot.notifier import FeishuNotifier, sanitize_and_shorten from community_qa_bot.text import split_text from community_qa_bot.types import FeedbackType, Platform logger = logging.getLogger(__name__) qa_service = CommunityQAService( api_key=TIDB_AI_API_KEY, chat_engine=TIDB_AI_DISCORD_CHAT_ENGINE, ) new_qa_service = CommunityNewQaService() bot = interactions.Client(intents=Intents.ALL, sync_interactions=True, asyncio_debug=True) @bot.listen() async def on_startup(): print(f"Bot({bot.user.id}) is ready!") def pure_mention_text(text: str) -> str: return text.replace(f'<@{bot.user.id}>', '').strip() def build_qa_context_for_thread(history: list[Message]) -> tuple[list, bool]: qa_context = [] for m in history: if m.type == MessageType.THREAD_STARTER_MESSAGE: # this is a trick to make the message reference work # to fetch the referenced message, we need to use the reference message's channel_id m._channel_id = m.message_reference.channel_id qa_context.append({ 'role': 'user', 'content': pure_mention_text(m.get_referenced_message().content) }) elif m.type == MessageType.DEFAULT: if m.author == bot.user: qa_context.append({ 'role': 'assistant', 'content': m.content }) elif m.contains_mention(bot.user): qa_context.append({ 'role': 'user', 'content': pure_mention_text(m.content) }) else: continue qa_context = qa_context[-5:] should_skip = False if not qa_context or qa_context[-1]['role'] == 'assistant': should_skip = True return qa_context, should_skip async def get_user(user_id: int, username: str, display_name: str) -> User: async with DBSession() as session: stmt = select(User).filter(User.id == user_id) result = await session.execute(stmt) user = result.scalars().first() if not user: user = User( id=user_id, username=username, display_name=display_name, user_type=Platform.DISCORD.value ) session.add(user) await session.commit() return user async def reach_daily_limit(user: User) -> bool: today = datetime.datetime.now().strftime('%Y-%m-%d') async with DBSession() as session: stmt = select(func.count()).where( and_( QaHistory.user_id == user.id, func.date(QaHistory.created_at) == today ) ) result = await session.execute(stmt) count = result.scalar_one() return count >= DISCORD_QA_DAILY_LIMIT async def create_qa_history(): async with DBSession() as session: QaHistory() @bot.listen() async def on_message_create(event: MessageCreate): message: "Message" = event.message channel = message.channel user = await get_user(message.author.id, message.author.tag, message.author.username) if message.author == bot.user: return if channel.type == ChannelType.DM: await channel.send(DISCORD_HELP_TEXT) return if not user.is_staff and \ message.type == MessageType.DEFAULT and \ channel.type in (ChannelType.GUILD_TEXT, ChannelType.GUILD_PUBLIC_THREAD): if (channel.type == ChannelType.GUILD_TEXT and channel.id in DISCORD_NEW_HELP_MESSAGE_ALERT_CHANNEL_IDS) or \ (channel.type == ChannelType.GUILD_PUBLIC_THREAD and message.thread and channel.parent_id in DISCORD_NEW_HELP_MESSAGE_ALERT_CHANNEL_IDS): shorten_content = sanitize_and_shorten(message.content) feishu_notifier = FeishuNotifier(DISCORD_NEW_HELP_MESSAGE_FEISHU_WEBHOOK) await asyncio.to_thread( feishu_notifier.send_markdown, title='New help message', content=f'Platform: Discord\nAuthor: {user.username}\n'\ f'Content: {shorten_content}\nLink: {message.jump_url}' ) if channel.type == ChannelType.GUILD_TEXT: if channel.id not in DISCORD_HELP_CHANNEL_IDS: logger.info(f"Guild Text Message({message.id}) is not in help channel, skip") return elif channel.type == ChannelType.GUILD_PUBLIC_THREAD: if channel.parent_id not in DISCORD_HELP_CHANNEL_IDS: logger.info(f"Guild Public Thread Message({message.id}) is not in help channel, skip") return else: return if not message.contains_mention(bot.user): return thread = channel if channel.type == ChannelType.GUILD_TEXT: question = pure_mention_text(message.content) thread = await message.create_thread(name=question[:100], auto_archive_duration=AutoArchiveDuration.ONE_HOUR) history = [message] elif channel.type == ChannelType.GUILD_PUBLIC_THREAD: history = await channel.history(limit=5, before=message.id).fetch() # message in history is in ascending order, and does not include the message itself history = list(reversed(history)) + [message] else: return if await reach_daily_limit(user): await thread.send(DISCORD_REACH_LIMIT_TEXT) return qa_context, should_skip = build_qa_context_for_thread(history) if should_skip: return async with thread.typing: start = time.time() # if channel.id in DISCORD_HELP_CHANNEL_NEW_QA_SERVICE_IDS or channel.parent_id in DISCORD_HELP_CHANNEL_NEW_QA_SERVICE_IDS: answer, ai_chat_id, ai_chat_message_id = await qa_service.ask_async(origin="discord", qa_context=qa_context) # else: # answer, ai_chat_id, ai_chat_message_id = await qa_service.ask_async(origin="discord", qa_context=qa_context) spend_time = time.time() - start if answer: # discord message max length is 2000 texts = split_text(answer, 2000) for i, text in enumerate(texts): if i == len(texts) - 1: components = [ ActionRow( Button( style=ButtonStyle.SECONDARY, emoji="👍", custom_id='qa_like' ), Button( style=ButtonStyle.SECONDARY, emoji="👎", custom_id='qa_dislike' ) ) ] else: components = None response_message = await thread.send(text, components=components) await asyncio.sleep(0.2) async with DBSession() as session: session.add(QaHistory( user_id=user.id, discord_message_id=message.id, discord_response_message_id=response_message.id, discord_channel_id=message.channel.id, request_body=qa_context, response_text=answer, ai_chat_id=ai_chat_id, ai_chat_message_id=ai_chat_message_id, spend_time=spend_time )) await session.commit() @bot.listen() async def on_component(event: Component): ctx = event.ctx user = await get_user(ctx.user.id, ctx.user.tag, ctx.user.username) async with DBSession() as session: stmt = select(QaHistory).filter( QaHistory.discord_response_message_id == ctx.message_id, ) result = await session.execute(stmt) qa_history = result.scalars().first() if qa_history: session.add(QaFeedback( qa_history_id=qa_history.id, user_id=user.id, feedback_type=FeedbackType.DISLIKE.value if ctx.custom_id == 'qa_dislike' else FeedbackType.LIKE.value, )) await session.commit() match ctx.custom_id: case 'qa_dislike': await qa_service.dislike_async( origin="discord", chat_message_id=qa_history.ai_chat_message_id, username=user.username, comment=f'From discord, user {user.username} click dislike button' ) await ctx.send( f'Thanks for your feedback, could you please give me some advice, {qa_history.get_signature_url(user.id)}', ephemeral=True ) case 'qa_like': await qa_service.like_async( origin="discord", chat_message_id=qa_history.ai_chat_message_id, username=user.username ) await ctx.send('Thanks for your feedback!', ephemeral=True) def run(): bot.load_extension('community_qa_bot.discord.commands') if DEBUG: # hot reload bot.load_extension("interactions.ext.jurigged") bot.start(DISCORD_BOT_TOKEN)