Module atomic_agents.lib.components.agent_memory

Classes

class AgentMemory (max_messages: Optional[int] = None)

Manages the chat history for an AI agent.

Attributes

history : List[Message]

A list of messages representing the chat history.

max_messages : Optional[int]

Maximum number of messages to keep in history.

current_turn_id : Optional[str]

The ID of the current turn.

Initializes the AgentMemory with an empty history and optional constraints.

Args

max_messages : Optional[int]

Maximum number of messages to keep in history. When exceeded, oldest messages are removed first.

Expand source code

class AgentMemory:
    """
    Manages the chat history for an AI agent.

    Attributes:
        history (List[Message]): A list of messages representing the chat history.
        max_messages (Optional[int]): Maximum number of messages to keep in history.
        current_turn_id (Optional[str]): The ID of the current turn.
    """

    def __init__(self, max_messages: Optional[int] = None):
        """
        Initializes the AgentMemory with an empty history and optional constraints.

        Args:
            max_messages (Optional[int]): Maximum number of messages to keep in history.
                When exceeded, oldest messages are removed first.
        """
        self.history: List[Message] = []
        self.max_messages = max_messages
        self.current_turn_id: Optional[str] = None

    def initialize_turn(self) -> None:
        """
        Initializes a new turn by generating a random turn ID.
        """
        self.current_turn_id = str(uuid.uuid4())

    def add_message(
        self,
        role: str,
        content: BaseIOSchema,
    ) -> None:
        """
        Adds a message to the chat history and manages overflow.

        Args:
            role (str): The role of the message sender.
            content (BaseIOSchema): The content of the message.
        """
        if self.current_turn_id is None:
            self.initialize_turn()

        message = Message(role=role, content=content, turn_id=self.current_turn_id)
        self.history.append(message)
        self._manage_overflow()

    def _manage_overflow(self) -> None:
        """
        Manages the chat history overflow based on max_messages constraint.
        """
        if self.max_messages is not None:
            while len(self.history) > self.max_messages:
                self.history.pop(0)

    def get_history(self) -> List[Dict]:
        """
        Retrieves the chat history, filtering out unnecessary fields and serializing content.

        Returns:
            List[Dict]: The list of messages in the chat history as dictionaries.
        """
        return [
            {
                "role": message.role,
                "content": json.dumps(message.content.model_dump()),
            }
            for message in self.history
        ]

    def copy(self) -> "AgentMemory":
        """
        Creates a copy of the chat memory.

        Returns:
            AgentMemory: A copy of the chat memory.
        """
        new_memory = AgentMemory(max_messages=self.max_messages)
        new_memory.load(self.dump())
        new_memory.current_turn_id = self.current_turn_id
        return new_memory

    def get_current_turn_id(self) -> Optional[str]:
        """
        Returns the current turn ID.

        Returns:
            Optional[str]: The current turn ID, or None if not set.
        """
        return self.current_turn_id

    def delete_turn_id(self, turn_id: int):
        """
        Delete messages from the memory by its turn ID.

        Args:
            turn_id (int): The turn ID of the message to delete.

        Returns:
            str: A success message with the deleted turn ID.

        Raises:
            ValueError: If the specified turn ID is not found in the memory.
        """
        initial_length = len(self.history)
        self.history = [msg for msg in self.history if msg.turn_id != turn_id]

        if len(self.history) == initial_length:
            raise ValueError(f"Turn ID {turn_id} not found in memory.")

        # Update current_turn_id if necessary
        if not len(self.history):
            self.current_turn_id = None
        elif turn_id == self.current_turn_id:
            # Always update to the last message's turn_id
            self.current_turn_id = self.history[-1].turn_id

    def get_message_count(self) -> int:
        """
        Returns the number of messages in the chat history.

        Returns:
            int: The number of messages.
        """
        return len(self.history)

    def dump(self) -> str:
        """
        Serializes the entire AgentMemory instance to a JSON string.

        Returns:
            str: A JSON string representation of the AgentMemory.
        """
        serialized_history = []
        for message in self.history:
            content_class = message.content.__class__
            serialized_message = {
                "role": message.role,
                "content": {
                    "class_name": f"{content_class.__module__}.{content_class.__name__}",
                    "data": message.content.model_dump(),
                },
                "turn_id": message.turn_id,
            }
            serialized_history.append(serialized_message)

        memory_data = {
            "history": serialized_history,
            "max_messages": self.max_messages,
            "current_turn_id": self.current_turn_id,
        }
        return json.dumps(memory_data)

    def load(self, serialized_data: str) -> None:
        """
        Deserializes a JSON string and loads it into the AgentMemory instance.

        Args:
            serialized_data (str): A JSON string representation of the AgentMemory.

        Raises:
            ValueError: If the serialized data is invalid or cannot be deserialized.
        """
        try:
            memory_data = json.loads(serialized_data)
            self.history = []
            self.max_messages = memory_data["max_messages"]
            self.current_turn_id = memory_data["current_turn_id"]

            for message_data in memory_data["history"]:
                content_info = message_data["content"]
                content_class = self._get_class_from_string(content_info["class_name"])
                content_instance = content_class(**content_info["data"])

                message = Message(role=message_data["role"], content=content_instance, turn_id=message_data["turn_id"])
                self.history.append(message)
        except (json.JSONDecodeError, KeyError, AttributeError, TypeError) as e:
            raise ValueError(f"Invalid serialized data: {e}")

    @staticmethod
    def _get_class_from_string(class_string: str) -> Type[BaseIOSchema]:
        """
        Retrieves a class object from its string representation.

        Args:
            class_string (str): The fully qualified class name.

        Returns:
            Type[BaseIOSchema]: The class object.

        Raises:
            AttributeError: If the class cannot be found.
        """
        module_name, class_name = class_string.rsplit(".", 1)
        module = __import__(module_name, fromlist=[class_name])
        return getattr(module, class_name)

Methods

def add_message(self, role: str, content: BaseIOSchema) ‑> None

Adds a message to the chat history and manages overflow.

Args

role : str

The role of the message sender.

content : BaseIOSchema

The content of the message.

def copy(self) ‑> AgentMemory

Creates a copy of the chat memory.

Returns

AgentMemory

A copy of the chat memory.

def delete_turn_id(self, turn_id: int)

Delete messages from the memory by its turn ID.

Args

turn_id : int

The turn ID of the message to delete.

Returns

str

A success message with the deleted turn ID.

Raises

ValueError

If the specified turn ID is not found in the memory.

def dump(self) ‑> str

Serializes the entire AgentMemory instance to a JSON string.

Returns

str

A JSON string representation of the AgentMemory.

def get_current_turn_id(self) ‑> Optional[str]

Returns the current turn ID.

Returns

Optional[str]

The current turn ID, or None if not set.

def get_history(self) ‑> List[Dict]

Retrieves the chat history, filtering out unnecessary fields and serializing content.

Returns

List[Dict]

The list of messages in the chat history as dictionaries.

def get_message_count(self) ‑> int

Returns the number of messages in the chat history.

Returns

int

The number of messages.

def initialize_turn(self) ‑> None

Initializes a new turn by generating a random turn ID.

def load(self, serialized_data: str) ‑> None

Deserializes a JSON string and loads it into the AgentMemory instance.

Args

serialized_data : str

A JSON string representation of the AgentMemory.

Raises

ValueError

If the serialized data is invalid or cannot be deserialized.

class Message (**data: Any)

Represents a message in the chat history.

Attributes

role : str

The role of the message sender (e.g., 'user', 'system', 'tool').

content : BaseIOSchema

The content of the message.

turn_id : Optional[str]

Unique identifier for the turn this message belongs to.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class Message(BaseModel):
    """
    Represents a message in the chat history.

    Attributes:
        role (str): The role of the message sender (e.g., 'user', 'system', 'tool').
        content (BaseIOSchema): The content of the message.
        turn_id (Optional[str]): Unique identifier for the turn this message belongs to.
    """

    role: str
    content: BaseIOSchema
    turn_id: Optional[str] = None

Ancestors

• pydantic.main.BaseModel

Class variables

var content : BaseIOSchema

The type of the None singleton.

var model_config

The type of the None singleton.

var role : str

The type of the None singleton.

var turn_id : Optional[str]

The type of the None singleton.