diff --git a/.vscode/settings.json b/.vscode/settings.json index 01e4dac..9adce29 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -23,17 +23,21 @@ "levelname", "lovibot", "Lovinator", + "Messageable", + "mountpoint", "ndarray", "nobot", "nparr", "numpy", "opencv", + "percpu", "plubplub", "pycodestyle", "pydocstyle", "pyproject", "PYTHONDONTWRITEBYTECODE", "PYTHONUNBUFFERED", + "Slowmode", "testpaths", "thelovinator", "tobytes", diff --git a/main.py b/main.py index 5667865..cba205c 100644 --- a/main.py +++ b/main.py @@ -84,7 +84,14 @@ class LoviBotClient(discord.Client): async with message.channel.typing(): try: - response: str | None = chat(incoming_message, openai_client, str(message.channel.id)) + response: str | None = chat( + user_message=incoming_message, + openai_client=openai_client, + current_channel=message.channel, + user=message.author, + allowed_users=allowed_users, + all_channels_in_guild=message.guild.channels if message.guild else None, + ) except openai.OpenAIError as e: logger.exception("An error occurred while chatting with the AI model.") e.add_note(f"Message: {incoming_message}\nEvent: {message}\nWho: {message.author.name}") @@ -92,8 +99,6 @@ class LoviBotClient(discord.Client): return if response: - response = f"{message.author.name}: {message.content}\n\n{response}" - logger.info("Responding to message: %s with: %s", incoming_message, response) await message.channel.send(response) @@ -173,7 +178,14 @@ async def ask(interaction: discord.Interaction, text: str) -> None: return try: - response: str | None = chat(text, openai_client, str(interaction.channel_id)) + response: str | None = chat( + user_message=text, + openai_client=openai_client, + current_channel=interaction.channel, + user=interaction.user, + allowed_users=allowed_users, + all_channels_in_guild=interaction.guild.channels if interaction.guild else None, + ) except openai.OpenAIError as e: logger.exception("An error occurred while chatting with the AI model.") await interaction.followup.send(f"An error occurred: {e}") @@ -375,9 +387,7 @@ def extract_image_url(message: discord.Message) -> str | None: break if not image_url: - match: re.Match[str] | None = re.search( - r"(https?://[^\s]+\.(png|jpg|jpeg|gif|webp)(\?[^\s]*)?)", message.content, re.IGNORECASE - ) + match: re.Match[str] | None = re.search(r"(https?://[^\s]+\.(png|jpg|jpeg|gif|webp)(\?[^\s]*)?)", message.content, re.IGNORECASE) if match: image_url = match.group(0) diff --git a/misc.py b/misc.py index 60c5a91..9633ea8 100644 --- a/misc.py +++ b/misc.py @@ -5,7 +5,15 @@ import logging from collections import deque from typing import TYPE_CHECKING +import psutil +from discord import Member, User, channel + if TYPE_CHECKING: + from collections.abc import Sequence + + from discord.abc import MessageableChannel + from discord.guild import GuildChannel + from discord.interactions import InteractionChannel from openai import OpenAI from openai.types.chat.chat_completion import ChatCompletion @@ -49,7 +57,7 @@ def add_message_to_memory(channel_id: str, user: str, message: str) -> None: logger.info("Added message to memory: %s from %s in channel %s", message, user, channel_id) -def get_recent_messages(channel_id: str, threshold_minutes: int = 10) -> list[tuple[str, str]]: +def get_recent_messages(channel_id: int, threshold_minutes: int = 10) -> list[tuple[str, str]]: """Retrieve messages from the last `threshold_minutes` minutes for a specific channel. Args: @@ -59,33 +67,131 @@ def get_recent_messages(channel_id: str, threshold_minutes: int = 10) -> list[tu Returns: A list of tuples containing user and message content. """ - if channel_id not in recent_messages: + if str(channel_id) not in recent_messages: return [] - threshold: datetime.datetime = datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta( - minutes=threshold_minutes - ) - return [(user, message) for user, message, timestamp in recent_messages[channel_id] if timestamp > threshold] + threshold: datetime.datetime = datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta(minutes=threshold_minutes) + return [(user, message) for user, message, timestamp in recent_messages[str(channel_id)] if timestamp > threshold] -def chat(user_message: str, openai_client: OpenAI, channel_id: str) -> str | None: +def extra_context(current_channel: MessageableChannel | InteractionChannel | None, user: User | Member) -> str: + """Add extra context to the chat prompt. + + For example: + - Current date and time + - Channel name and server + - User's current status (online/offline) + - User's role in the server (e.g., admin, member) + - CPU usage + - Memory usage + - Disk usage + - How many messages saved in memory + + Args: + current_channel: The channel where the conversation is happening. + user: The user who is interacting with the bot. + + Returns: + The extra context to include in the chat prompt. + """ + context: str = "" + + # Information about the servers and channels: + context += "KillYoy's Server Information:\n" + context += "- Server is for friends to hang out and chat.\n" + context += "- Server was created by KillYoy (<@98468214824001536>)\n" + + # Current date and time + context += f"Current date and time: {datetime.datetime.now(tz=datetime.UTC)} UTC, but user is in CEST or CET\n" + + # Channel name and server + if isinstance(current_channel, channel.TextChannel): + context += f"Channel name: {current_channel.name}, channel ID: {current_channel.id}, Server: {current_channel.guild.name}\n" + + # User information + context += f"User name: {user.name}, User ID: {user.id}\n" + if isinstance(user, Member): + context += f"User roles: {', '.join([role.name for role in user.roles])}\n" + context += f"User status: {user.status}\n" + context += f"User is currently {'on mobile' if user.is_on_mobile() else 'on desktop'}\n" + context += f"User joined server at: {user.joined_at}\n" + context += f"User's current activity: {user.activity}\n" + context += f"User's username color: {user.color}\n" + + # System information + context += f"CPU usage per core: {psutil.cpu_percent(percpu=True)}%\n" + context += f"Memory usage: {psutil.virtual_memory().percent}%\n" + context += f"Total memory: {psutil.virtual_memory().total / (1024 * 1024):.2f} MB\n" + context += f"Swap memory usage: {psutil.swap_memory().percent}%\n" + context += f"Swap memory total: {psutil.swap_memory().total / (1024 * 1024):.2f} MB\n" + context += f"Bot memory usage: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB\n" + uptime: datetime.timedelta = datetime.datetime.now(tz=datetime.UTC) - datetime.datetime.fromtimestamp(psutil.boot_time(), tz=datetime.UTC) + context += f"System uptime: {uptime}\n" + context += "Disk usage:\n" + for partition in psutil.disk_partitions(): + try: + context += f" {partition.mountpoint}: {psutil.disk_usage(partition.mountpoint).percent}%\n" + except PermissionError as e: + context += f" {partition.mountpoint} got PermissionError: {e}\n" + + if current_channel: + context += f"Messages saved in memory: {len(get_recent_messages(channel_id=current_channel.id))}\n" + + return context + + +def chat( # noqa: PLR0913, PLR0917 + user_message: str, + openai_client: OpenAI, + current_channel: MessageableChannel | InteractionChannel | None, + user: User | Member, + allowed_users: list[str], + all_channels_in_guild: Sequence[GuildChannel] | None = None, +) -> str | None: """Chat with the bot using the OpenAI API. Args: user_message: The message to send to OpenAI. openai_client: The OpenAI client to use. - channel_id: The ID of the channel where the conversation is happening. + current_channel: The channel where the conversation is happening. + user: The user who is interacting with the bot. + allowed_users: The list of allowed users to interact with the bot. + all_channels_in_guild: The list of all channels in the guild. Returns: The response from the AI model. """ - # Include recent messages in the prompt - recent_context: str = "\n".join([f"{user}: {message}" for user, message in get_recent_messages(channel_id)]) + recent_context: str = "" + context: str = "" + + if current_channel: + channel_id = int(current_channel.id) + recent_context: str = "\n".join([f"{user}: {message}" for user, message in get_recent_messages(channel_id=channel_id)]) + + context = extra_context(current_channel=current_channel, user=user) + + context += "The bot is in the following channels:\n" + if all_channels_in_guild: + for c in all_channels_in_guild: + context += f"{c!r}\n" + + context += "\nThe bot responds to the following users:\n" + for user_id in allowed_users: + context += f" - User ID: {user_id}\n" + prompt: str = ( - "You are in a Discord group chat. People can ask you questions. " + "You are in a Discord group chat. People can ask you questions.\n" "Use Discord Markdown to format messages if needed.\n" - f"Recent context:\n{recent_context}\n" + "Don't use emojis.\n" + "Extra context starts here:\n" + f"{context}" + "Extra context ends here.\n" + "Recent context starts here:\n" + f"{recent_context}\n" + "Recent context ends here.\n" + "User message starts here:\n" f"User: {user_message}" + "User message ends here.\n" ) completion: ChatCompletion = openai_client.chat.completions.create( diff --git a/pyproject.toml b/pyproject.toml index ee95242..ae17334 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,9 +7,11 @@ requires-python = ">=3.13" dependencies = [ "audioop-lts", "discord-py", + "httpx>=0.28.1", "numpy", "openai", "opencv-contrib-python-headless", + "psutil>=7.0.0", "python-dotenv", "sentry-sdk", ] @@ -26,7 +28,6 @@ lint.fixable = ["ALL"] lint.pydocstyle.convention = "google" lint.isort.required-imports = ["from __future__ import annotations"] lint.pycodestyle.ignore-overlong-task-comments = true -line-length = 120 lint.ignore = [ "CPY001", # Checks for the absence of copyright notices within Python files. @@ -53,6 +54,7 @@ lint.ignore = [ "Q003", # Checks for strings that include escaped quotes, and suggests changing the quote style to avoid the need to escape them. "W191", # Checks for indentation that uses tabs. ] +line-length = 160 [tool.ruff.format]