Enhance chat functionality by adding extra context and improving message handling

This commit is contained in:
Joakim Hellsén 2025-08-17 03:57:41 +02:00
commit f0f4e3c9b7
4 changed files with 142 additions and 20 deletions

View file

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

24
main.py
View file

@ -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)

130
misc.py
View file

@ -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(

View file

@ -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]