app/core/memory.py

from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from app.config import config
from app.utils.logger import logger


class ConversationMemory:
    def __init__(self):
        self.max_messages = config["rag"]["memory"]["max_messages"]
        self.redis_url = config["database"]["redis"]["url"]
    
    def get_message_history(self, session_id: str) -> RedisChatMessageHistory:
        return RedisChatMessageHistory(
            session_id=session_id,
            url=self.redis_url,
            ttl=86400
        )
    
    def create_history_runnable(self, runnable):
        return RunnableWithMessageHistory(
            runnable,
            self.get_message_history,
            input_messages_key="input",
            history_messages_key="chat_history"
        )
    
    def get_messages(self, session_id: str):
        history = self.get_message_history(session_id)
        messages = history.messages
        return messages[-self.max_messages:] if len(messages) > self.max_messages else messages
    
    def add_message(self, session_id: str, role: str, content: str):
        """Add a message to the conversation history."""
        try:
            history = self.get_message_history(session_id)
            if role == "user":
                history.add_user_message(content)
            elif role == "assistant":
                history.add_ai_message(content)
            logger.debug(f"Added {role} message to session {session_id}")
        except Exception as e:
            logger.error(f"Error adding message: {e}")
    
    def clear(self, session_id: str):
        try:
            history = self.get_message_history(session_id)
            history.clear()
            logger.info(f"Cleared memory for session {session_id}")
        except Exception as e:
            logger.error(f"Memory clear error: {e}")


conversation_memory = ConversationMemory()