memory.py

"""
This module defines a `ConversationMemoryRunnable` class that stores and manages
conversation history for an LLM-based assistant.
"""

from langchain_core.runnables import Runnable

# Initialize a global variable to store the conversation history
conversation_history = []

class ConversationMemoryRunnable(Runnable):
    """
    A runnable that manages conversation memory by storing user inputs and
    assistant responses. It uses an LLM to generate responses based on session history.
    """

    def __init__(self, llm):
        """
        Initializes the ConversationMemoryRunnable with an LLM.

        Args:
            llm: The language model used to generate responses.
        """
        self.llm = llm

    def run(self, input_text: str, **kwargs) -> str:
        """
        Generates a response from the LLM based on conversation history and
        the latest user input.

        Args:
            input_text: The user's input message.
            **kwargs: Additional keyword arguments for the LLM.

        Returns:
            str: The response generated by the LLM.
        """
        session_history = get_session_history()
        combined_input = f"{session_history}\nUser: {input_text}\nAssistant:"
        response = self.llm.invoke(combined_input, **kwargs)
        update_conversation_history(input_text, response)
        return response

    def invoke(self, *args, **kwargs) -> str:
        """
        Invokes the run method using the first argument as input text.

        Args:
            *args: Positional arguments where the first argument is expected to be the input text.
            **kwargs: Additional keyword arguments for the run method.

        Returns:
            str: The response generated by the run method.
        """
        input_text = args[0] if args else ""
        return self.run(input_text, **kwargs)

def get_session_history() -> str:
    """
    Retrieves the formatted session history of conversation between the user and the assistant.

    Returns:
        str: A formatted string of the conversation history.
    """
    formatted_history = "\n".join(
        f"User: {entry['user']}\nAssistant: {entry['assistant']}" for entry in conversation_history
    )
    return formatted_history

def update_conversation_history(user_message: str, assistant_response: str) -> None:
    """
    Updates the conversation history by adding the latest user message and assistant response.

    Args:
        user_message: The user's input message.
        assistant_response: The assistant's response message.
    """
    conversation_history.append({"user": user_message, "assistant": assistant_response})
    if len(conversation_history) > 2:  # Adjust the limit as needed
        conversation_history.pop(0)  # Remove the oldest entry