Hiddenden/loyal_companion
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add PostgreSQL memory system for persistent user and conversation storage

Browse Source 
- Add PostgreSQL database with SQLAlchemy async support
- Create models: User, UserFact, UserPreference, Conversation, Message, Guild, GuildMember
- Add custom name support so bot knows 'who is who'
- Add user facts system for remembering things about users
- Add persistent conversation history that survives restarts
- Add memory commands cog (!setname, !remember, !whatdoyouknow, !forgetme)
- Add admin commands (!setusername, !teachbot)
- Set up Alembic for database migrations
- Update docker-compose with PostgreSQL service
- Gracefully falls back to in-memory storage when DB not configured
This commit is contained in:
Latte
2026-01-12 14:00:06 +01:00
parent 853e7c9fcd
commit e00d4fd501
20 changed files with 1623 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

98
src/daemon_boyfriend/cogs/ai_chat.py
View File

@@ -12,7 +12,10 @@ from daemon_boyfriend.services import (
ConversationManager,
ImageAttachment,
Message,
PersistentConversationManager,
SearXNGService,
UserService,
db,
)
from daemon_boyfriend.utils import get_monitor
@@ -77,11 +80,17 @@ class AIChatCog(commands.Cog):
def __init__(self, bot: commands.Bot) -> None:
self.bot = bot
self.ai_service = AIService()
# Fallback in-memory conversation manager (used when DB not configured)
self.conversations = ConversationManager()
self.search_service: SearXNGService | None = None
if settings.searxng_enabled and settings.searxng_url:
self.search_service = SearXNGService(settings.searxng_url)
@property
def use_database(self) -> bool:
"""Check if database is available for use."""
return db.is_initialized
@commands.Cog.listener()
async def on_message(self, message: discord.Message) -> None:
"""Respond when the bot is mentioned."""
@@ -395,6 +404,95 @@ class AIChatCog(commands.Cog):
Returns:
The AI's response text
"""
if self.use_database:
return await self._generate_response_with_db(message, user_message)
else:
return await self._generate_response_in_memory(message, user_message)
async def _generate_response_with_db(self, message: discord.Message, user_message: str) -> str:
"""Generate response using database-backed storage."""
async with db.session() as session:
user_service = UserService(session)
conv_manager = PersistentConversationManager(session)
# Get or create user
user = await user_service.get_or_create_user(
discord_id=message.author.id,
username=message.author.name,
display_name=message.author.display_name,
)
# Get or create conversation
conversation = await conv_manager.get_or_create_conversation(
user=user,
guild_id=message.guild.id if message.guild else None,
channel_id=message.channel.id,
)
# Get history
history = await conv_manager.get_history(conversation)
# Extract any image attachments from the message
images = self._extract_image_attachments(message)
image_urls = [img.url for img in images] if images else None
# Add current message to history for the API call
current_message = Message(role="user", content=user_message, images=images)
messages = history + [current_message]
# Check if we should search the web
search_context = await self._maybe_search(user_message)
# Get context about mentioned users
mentioned_users_context = self._get_mentioned_users_context(message)
# Build system prompt with additional context
system_prompt = self.ai_service.get_system_prompt()
# Add user context from database (custom name, known facts)
user_context = await user_service.get_user_context(user)
system_prompt += f"\n\n--- User Context ---\n{user_context}"
# Add mentioned users context
if mentioned_users_context:
system_prompt += f"\n\n--- {mentioned_users_context} ---"
# Add search results if available
if search_context:
system_prompt += (
"\n\n--- Web Search Results ---\n"
"Use the following current information from the web to help answer the user's question. "
"Cite sources when relevant.\n\n"
f"{search_context}"
)
# Generate response
response = await self.ai_service.chat(
messages=messages,
system_prompt=system_prompt,
)
# Save the exchange to database
await conv_manager.add_exchange(
conversation=conversation,
user=user,
user_message=user_message,
assistant_message=response.content,
discord_message_id=message.id,
image_urls=image_urls,
)
logger.debug(
f"Generated response for user {user.discord_id}: "
f"{len(response.content)} chars, {response.usage}"
)
return response.content
async def _generate_response_in_memory(
self, message: discord.Message, user_message: str
) -> str:
"""Generate response using in-memory storage (fallback)."""
user_id = message.author.id
# Get conversation history

260
src/daemon_boyfriend/cogs/memory.py Normal file
View File

@@ -0,0 +1,260 @@
"""Memory management cog - commands for managing bot memory about users."""
import logging
import discord
from discord.ext import commands
from daemon_boyfriend.services import UserService, db
logger = logging.getLogger(__name__)
class MemoryCog(commands.Cog):
"""Commands for managing bot memory about users."""
def __init__(self, bot: commands.Bot) -> None:
self.bot = bot
def _check_database(self) -> bool:
"""Check if database is available."""
return db.is_initialized
@commands.command(name="setname")
async def set_name(self, ctx: commands.Context, *, name: str) -> None:
"""Set your preferred name for the bot to use.
Usage: !setname John
"""
if not self._check_database():
await ctx.reply("Memory features are not available (database not configured).")
return
if len(name) > 100:
await ctx.reply("Name is too long! Please use 100 characters or less.")
return
async with db.session() as session:
user_service = UserService(session)
user = await user_service.get_or_create_user(
discord_id=ctx.author.id,
username=ctx.author.name,
display_name=ctx.author.display_name,
)
await user_service.set_custom_name(ctx.author.id, name)
await ctx.reply(f"Got it! I'll call you **{name}** from now on.")
@commands.command(name="clearname")
async def clear_name(self, ctx: commands.Context) -> None:
"""Clear your custom name and use your Discord name instead.
Usage: !clearname
"""
if not self._check_database():
await ctx.reply("Memory features are not available (database not configured).")
return
async with db.session() as session:
user_service = UserService(session)
await user_service.set_custom_name(ctx.author.id, None)
await ctx.reply("Done! I'll use your Discord display name now.")
@commands.command(name="remember")
async def remember_fact(self, ctx: commands.Context, *, fact: str) -> None:
"""Tell the bot something to remember about you.
Usage: !remember I love pizza
Usage: !remember My favorite color is blue
"""
if not self._check_database():
await ctx.reply("Memory features are not available (database not configured).")
return
if len(fact) > 500:
await ctx.reply("That's too long to remember! Please keep it under 500 characters.")
return
async with db.session() as session:
user_service = UserService(session)
user = await user_service.get_or_create_user(
discord_id=ctx.author.id,
username=ctx.author.name,
display_name=ctx.author.display_name,
)
await user_service.add_fact(
user=user,
fact_type="general",
fact_content=fact,
source="explicit",
confidence=1.0,
)
await ctx.reply(f"I'll remember that!")
@commands.command(name="whatdoyouknow", aliases=["aboutme", "myinfo"])
async def what_do_you_know(self, ctx: commands.Context) -> None:
"""Show what the bot remembers about you.
Usage: !whatdoyouknow
"""
if not self._check_database():
await ctx.reply("Memory features are not available (database not configured).")
return
async with db.session() as session:
user_service = UserService(session)
user = await user_service.get_user_by_discord_id(ctx.author.id)
if not user:
await ctx.reply("I don't have any information about you yet!")
return
facts = await user_service.get_user_facts(user, active_only=True)
embed = discord.Embed(
title=f"What I know about {user.display_name}",
color=discord.Color.blue(),
)
embed.add_field(
name="Discord Username",
value=user.discord_username,
inline=True,
)
if user.custom_name:
embed.add_field(
name="Preferred Name",
value=user.custom_name,
inline=True,
)
embed.add_field(
name="First Seen",
value=user.first_seen_at.strftime("%Y-%m-%d"),
inline=True,
)
if facts:
facts_text = "\n".join(f"- {fact.fact_content}" for fact in facts[:15])
if len(facts) > 15:
facts_text += f"\n... and {len(facts) - 15} more"
embed.add_field(
name=f"Things I Remember ({len(facts)})",
value=facts_text or "Nothing yet!",
inline=False,
)
else:
embed.add_field(
name="Things I Remember",
value="Nothing yet! Use `!remember` to tell me something.",
inline=False,
)
await ctx.reply(embed=embed)
@commands.command(name="forgetme")
async def forget_me(self, ctx: commands.Context) -> None:
"""Clear all facts the bot knows about you.
Usage: !forgetme
"""
if not self._check_database():
await ctx.reply("Memory features are not available (database not configured).")
return
async with db.session() as session:
user_service = UserService(session)
user = await user_service.get_user_by_discord_id(ctx.author.id)
if not user:
await ctx.reply("I don't have any information about you to forget!")
return
count = await user_service.delete_user_facts(user)
if count > 0:
await ctx.reply(f"Done! I've forgotten {count} thing(s) about you.")
else:
await ctx.reply("I didn't have anything to forget about you!")
@commands.command(name="setusername")
@commands.has_permissions(administrator=True)
async def set_user_name(
self, ctx: commands.Context, user: discord.Member, *, name: str
) -> None:
"""[Admin] Set a custom name for another user.
Usage: !setusername @user John
"""
if not self._check_database():
await ctx.reply("Memory features are not available (database not configured).")
return
if len(name) > 100:
await ctx.reply("Name is too long! Please use 100 characters or less.")
return
async with db.session() as session:
user_service = UserService(session)
db_user = await user_service.get_or_create_user(
discord_id=user.id,
username=user.name,
display_name=user.display_name,
)
await user_service.set_custom_name(user.id, name)
await ctx.reply(f"Got it! I'll call {user.mention} **{name}** from now on.")
@commands.command(name="teachbot")
@commands.has_permissions(administrator=True)
async def teach_bot(self, ctx: commands.Context, user: discord.Member, *, fact: str) -> None:
"""[Admin] Teach the bot a fact about a user.
Usage: !teachbot @user They are a software developer
"""
if not self._check_database():
await ctx.reply("Memory features are not available (database not configured).")
return
if len(fact) > 500:
await ctx.reply("That's too long! Please keep it under 500 characters.")
return
async with db.session() as session:
user_service = UserService(session)
db_user = await user_service.get_or_create_user(
discord_id=user.id,
username=user.name,
display_name=user.display_name,
)
await user_service.add_fact(
user=db_user,
fact_type="general",
fact_content=fact,
source="admin",
confidence=1.0,
)
await ctx.reply(f"I'll remember that about {user.mention}!")
@set_user_name.error
@teach_bot.error
async def admin_command_error(
self, ctx: commands.Context, error: commands.CommandError
) -> None:
"""Handle errors for admin commands."""
if isinstance(error, commands.MissingPermissions):
await ctx.reply("You need administrator permissions to use this command.")
elif isinstance(error, commands.MemberNotFound):
await ctx.reply("I couldn't find that user.")
else:
logger.error(f"Error in admin command: {error}", exc_info=True)
await ctx.reply(f"An error occurred: {error}")
async def setup(bot: commands.Bot) -> None:
"""Load the Memory cog."""
await bot.add_cog(MemoryCog(bot))