Enhance chat functionality by adding extra context and improving message handling

.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",

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(
+        match: re.Match[str] | None = re.search(r"(https?://[^\s]+\.(png|jpg|jpeg|gif|webp)(\?[^\s]*)?)", message.content, re.IGNORECASE)
-            r"(https?://[^\s]+\.(png|jpg|jpeg|gif|webp)(\?[^\s]*)?)", message.content, re.IGNORECASE
-        )
         if match:
             image_url = match.group(0)
 

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(
+    threshold: datetime.datetime = datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta(minutes=threshold_minutes)
-        minutes=threshold_minutes
+    return [(user, message) for user, message, timestamp in recent_messages[str(channel_id)] if timestamp > threshold]
-    )
-    return [(user, message) for user, message, timestamp in recent_messages[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 = ""
-    recent_context: str = "\n".join([f"{user}: {message}" for user, message in get_recent_messages(channel_id)])
+    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(

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]