KaniumCogs/reginaldCog/memory.py

import datetime
import random
import re
from collections import Counter

import discord
import openai
from openai import OpenAIError
from redbot.core import commands


class MemoryMixin:
    """Handles all memory-related functions for Reginald."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.short_term_memory_limit = 50
        self.summary_retention_limit = 25
        self.summary_retention_ratio = 0.8

    @commands.command(name="reginald_clear_short", help="Clears short-term memory for this channel.")
    @commands.has_permissions(administrator=True)
    async def clear_short_memory(self, ctx):
        async with self.config.guild(ctx.guild).short_term_memory() as short_memory:
            short_memory[str(ctx.channel.id)] = []
        await ctx.send("Short-term memory for this channel has been cleared.")

    @commands.command(name="reginald_set_limit", help="Set the short-term memory message limit.")
    @commands.has_permissions(administrator=True)
    async def set_short_term_memory_limit(self, ctx, limit: int):
        if limit < 5:
            await ctx.send("The short-term memory limit must be at least 5.")
            return

        self.short_term_memory_limit = limit
        await ctx.send(f"Short-term memory limit set to {limit} messages.")

    @commands.command(name="reginald_memory_limit", help="Displays the current short-term memory message limit.")
    async def get_short_term_memory_limit(self, ctx):
        await ctx.send(f"Current short-term memory limit: {self.short_term_memory_limit} messages.")

    @commands.command(name="reginald_clear_mid", help="Clears mid-term memory (summarized logs).")
    @commands.has_permissions(administrator=True)
    async def clear_mid_memory(self, ctx):
        async with self.config.guild(ctx.guild).mid_term_memory() as mid_memory:
            mid_memory[str(ctx.channel.id)] = []
        await ctx.send("Mid-term memory for this channel has been cleared.")

    @commands.command(name="reginald_summary", help="Displays a selected mid-term summary for this channel.")
    async def get_mid_term_summary(self, ctx, index: int):
        async with self.config.guild(ctx.guild).mid_term_memory() as mid_memory:
            summaries = mid_memory.get(str(ctx.channel.id), [])
            if not isinstance(summaries, list):
                summaries = []

            if not summaries:
                await ctx.send("No summaries available for this channel.")
                return

            if index < 1 or index > len(summaries):
                await ctx.send(f"Invalid index. Please provide a number between 1 and {len(summaries)}.")
                return

            selected_summary = summaries[index - 1]
            formatted_summary = (
                f"Summary {index} of {len(summaries)}\n"
                f"Date: {selected_summary.get('timestamp', 'Unknown')}\n"
                f"Topics: {', '.join(selected_summary.get('topics', [])) or 'None'}\n"
                f"Summary:\n\n{selected_summary.get('summary', '')}"
            )

            await self.send_long_message(ctx, formatted_summary)

    @commands.command(name="reginald_summaries", help="Lists available summaries for this channel.")
    async def list_mid_term_summaries(self, ctx):
        async with self.config.guild(ctx.guild).mid_term_memory() as mid_memory:
            summaries = mid_memory.get(str(ctx.channel.id), [])
            if not isinstance(summaries, list):
                summaries = []

            if not summaries:
                await ctx.send("No summaries available for this channel.")
                return

            summary_list = "\n".join(
                f"{i + 1}. {entry.get('timestamp', 'Unknown')} | Topics: {', '.join(entry.get('topics', [])) or 'None'}"
                for i, entry in enumerate(summaries)
            )

            await ctx.send(f"Available summaries:\n{summary_list[:2000]}")

    @commands.command(name="reginald_clear_long", help="Clears all long-term stored knowledge.")
    @commands.has_permissions(administrator=True)
    async def clear_long_memory(self, ctx):
        async with self.config.guild(ctx.guild).long_term_profiles() as long_memory:
            long_memory.clear()
        await ctx.send("All long-term memory has been erased.")

    @commands.command(name="reginald_reset_all", help="Completely resets all memory.")
    @commands.has_permissions(administrator=True)
    async def reset_all_memory(self, ctx):
        async with self.config.guild(ctx.guild).short_term_memory() as short_memory:
            short_memory.clear()
        async with self.config.guild(ctx.guild).mid_term_memory() as mid_memory:
            mid_memory.clear()
        async with self.config.guild(ctx.guild).long_term_profiles() as long_memory:
            long_memory.clear()
        await ctx.send("All memory has been completely reset.")

    @commands.command(name="reginald_recall", help="Recalls what Reginald knows about a user.")
    async def recall_user(self, ctx, user: discord.User):
        async with self.config.guild(ctx.guild).long_term_profiles() as long_memory:
            profile = long_memory.get(str(user.id), {})
            facts = profile.get("facts", [])

        if facts:
            recall_lines = [f"- {fact.get('fact', '')}" for fact in facts[:10]]
            await ctx.send(f"Memory recall for {user.display_name}:\n" + "\n".join(recall_lines))
        else:
            await ctx.send(f"No stored information on {user.display_name}.")

    @commands.command(name="reginald_forget", help="Forgets a specific user's long-term profile.")
    @commands.has_permissions(administrator=True)
    async def forget_user(self, ctx, user: discord.User):
        async with self.config.guild(ctx.guild).long_term_profiles() as long_memory:
            if str(user.id) in long_memory:
                del long_memory[str(user.id)]
                await ctx.send(f"Reginald has forgotten all stored information about {user.display_name}.")
            else:
                await ctx.send(f"No stored knowledge about {user.display_name} to delete.")

    async def summarize_memory(self, ctx, messages):
        summary_prompt = (
            "Summarize the following conversation into a structured, concise format that retains key details while maximizing brevity. "
            "Organize into sections: Key Takeaways, Disputed Points, Notable User Contributions, and Additional Context. "
            "Avoid repetition while preserving essential meaning."
        )

        summary_text = "\n".join(f"{msg.get('user', 'Unknown')}: {msg.get('content', '')}" for msg in messages)

        try:
            api_key = await self.config.guild(ctx.guild).openai_api_key()
            if not api_key:
                return (
                    "It appears that I have not been furnished with the necessary credentials to carry out this task. "
                    "Might I suggest consulting an administrator to rectify this unfortunate oversight?"
                )

            model = await self.config.openai_model() or "gpt-4o-mini"
            client = openai.AsyncOpenAI(api_key=api_key)
            response = await client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": summary_prompt},
                    {"role": "user", "content": summary_text},
                ],
                max_tokens=2048,
            )

            summary_content = (response.choices[0].message.content or "").strip()
            if not summary_content:
                return (
                    "Ah, an unusual predicament indeed. It seems that my attempt at summarization has resulted in "
                    "a void of information. I shall endeavor to be more verbose next time."
                )

            return summary_content

        except OpenAIError as error:
            error_message = f"OpenAI Error: {error}"
            reginald_responses = [
                f"Regrettably, I must inform you that I have encountered a bureaucratic obstruction whilst attempting to summarize:\n\n{error_message}",
                f"It would seem that a most unfortunate technical hiccup has befallen my faculties in the matter of summarization:\n\n{error_message}",
                f"Ah, it appears I have received an urgent memorandum stating that my summarization efforts have been thwarted:\n\n{error_message}",
                f"I regret to inform you that my usual eloquence is presently obstructed by an unforeseen complication while summarizing:\n\n{error_message}",
            ]
            return random.choice(reginald_responses)

    def extract_topics_from_summary(self, summary):
        keywords = re.findall(r"\b\w+\b", summary.lower())
        word_counts = Counter(keywords)

        stop_words = {
            "the",
            "and",
            "of",
            "in",
            "to",
            "is",
            "on",
            "for",
            "with",
            "at",
            "by",
            "it",
            "this",
            "that",
            "his",
            "her",
        }
        filtered_words = {
            word: count
            for word, count in word_counts.items()
            if word not in stop_words and len(word) > 2
        }

        topics = sorted(filtered_words, key=filtered_words.get, reverse=True)[:5]
        return topics

    def select_relevant_summaries(self, summaries, prompt):
        summaries = [summary for summary in summaries if isinstance(summary, dict)]
        if not summaries:
            return []

        max_summaries = 5 if len(prompt) > 50 else 3
        current_time = datetime.datetime.now()

        def calculate_weight(summary):
            topics = summary.get("topics", [])
            topic_match = sum(1 for topic in topics if topic in prompt.lower())
            frequency_score = len(topics)

            try:
                timestamp = datetime.datetime.strptime(summary.get("timestamp", ""), "%Y-%m-%d %H:%M")
                recency_factor = max(0.1, 1 - ((current_time - timestamp).days / 365))
            except ValueError:
                recency_factor = 0.1

            return (topic_match * 2) + (frequency_score * 1.5) + (recency_factor * 3)

        weighted_summaries = sorted(summaries, key=calculate_weight, reverse=True)
        return weighted_summaries[:max_summaries]

    def extract_fact_from_response(self, response_text):
        fact_patterns = [
            r"I recall that you (.*?)\.",
            r"You once mentioned that you (.*?)\.",
            r"Ah, you previously stated that (.*?)\.",
            r"As I remember, you (.*?)\.",
            r"I believe you (.*?)\.",
            r"I seem to recall that you (.*?)\.",
            r"You have indicated in the past that you (.*?)\.",
            r"From what I remember, you (.*?)\.",
            r"You previously mentioned that (.*?)\.",
            r"It is my understanding that you (.*?)\.",
            r"If I am not mistaken, you (.*?)\.",
        ]

        for pattern in fact_patterns:
            match = re.search(pattern, response_text, re.IGNORECASE)
            if match:
                return match.group(1)

        return None

    @commands.command(name="reginald_memory_status", help="Displays a memory usage summary.")
    async def memory_status(self, ctx):
        async with self.config.guild(ctx.guild).short_term_memory() as short_memory, self.config.guild(
            ctx.guild
        ).mid_term_memory() as mid_memory, self.config.guild(ctx.guild).long_term_profiles() as long_memory:
            short_count = sum(len(v) for v in short_memory.values() if isinstance(v, list))
            mid_count = sum(len(v) for v in mid_memory.values() if isinstance(v, list))
            long_count = len(long_memory)

        status_message = (
            "Memory status:\n"
            f"- Short-term messages stored: {short_count}\n"
            f"- Mid-term summaries stored: {mid_count}\n"
            f"- Long-term profiles stored: {long_count}\n"
        )
        await ctx.send(status_message)

    def normalize_fact(self, fact: str) -> str:
        return re.sub(r"\s+", " ", fact.strip().lower())

    async def update_long_term_memory(self, ctx, user_id: str, fact: str, source_message: str, timestamp: str):
        fact = self.normalize_fact(fact)
        if not fact:
            return

        async with self.config.guild(ctx.guild).long_term_profiles() as long_memory:
            if user_id not in long_memory:
                long_memory[user_id] = {"facts": []}

            user_facts = long_memory[user_id]["facts"]

            for entry in user_facts:
                if self.normalize_fact(entry.get("fact", "")) == fact:
                    entry["last_updated"] = timestamp
                    return

            conflicting_entry = None
            for entry in user_facts:
                existing_keywords = set(entry.get("fact", "").lower().split())
                new_keywords = set(fact.lower().split())
                if len(existing_keywords & new_keywords) >= 2:
                    conflicting_entry = entry
                    break

            if conflicting_entry is not None:
                conflicting_entry.setdefault("previous_versions", [])
                conflicting_entry["previous_versions"].append(
                    {
                        "fact": conflicting_entry.get("fact", ""),
                        "source": conflicting_entry.get("source", ""),
                        "timestamp": conflicting_entry.get("timestamp", ""),
                    }
                )
                conflicting_entry["fact"] = fact
                conflicting_entry["source"] = source_message
                conflicting_entry["timestamp"] = timestamp
                conflicting_entry["last_updated"] = timestamp
            else:
                user_facts.append(
                    {
                        "fact": fact,
                        "source": source_message,
                        "timestamp": timestamp,
                        "last_updated": timestamp,
                        "previous_versions": [],
                    }
                )