Package com.tnsai.agents.execution

Class MemoryManager

java.lang.Object
com.tnsai.agents.execution.MemoryManager
public class MemoryManager extends Object
Manages agent memory and state checkpointing.

This class is extracted from Agent to provide focused memory management:

  • Conversation history management
  • State checkpointing and restoration
  • Token-based pruning
  • Checkpoint history with rollback capabilities

Checkpoint System

Checkpoints capture the full agent state at a point in time, allowing rollback to previous states. This is useful for:

  • Error recovery
  • State exploration (try different paths)
  • Debugging and testing

Usage


 MemoryManager controller = new MemoryManager(agentId, memoryStore);

 // Specure
 controller.setMaxCheckpoints(5);
 controller.setMaxContextTokens(4000);

 // Create checkpoint before risky operation
 AgentCheckpoint checkpoint = controller.captureCheckpoint("before_action");

 // If operation fails, rollback
 controller.restoreCheckpoint(checkpoint);

 // Or rollback to last checkpoint
 controller.rollbackToLast();
See Also:

  • Constructor Details

    • MemoryManager

      public MemoryManager(String agentId, MemoryStore memoryStore)
      Creates a new MemoryManager.
      Parameters:
      agentId - the agent ID
      memoryStore - the memory store to manage

  • Method Details

    • setMaxCheckpoints

      public MemoryManager setMaxCheckpoints(int maxCheckpoints)
      Sets the maximum number of checkpoints to retain. Oldest checkpoints are removed when limit is exceeded.
      Parameters:
      maxCheckpoints - maximum checkpoints (must be positive)
      Returns:
      this for chaining

    • setMaxContextTokens

      public MemoryManager setMaxContextTokens(Integer maxTokens)
      Sets the maximum context token limit. When set, history is automatically pruned after adding messages.
      Parameters:
      maxTokens - maximum tokens (null to disable)
      Returns:
      this for chaining

    • addUserMessage

      public void addUserMessage(String content)
      Adds a user message to the conversation history.
      Parameters:
      content - the message content

    • addAssistantMessage

      public void addAssistantMessage(String content)
      Adds an assistant message to the conversation history.
      Parameters:
      content - the message content

    • addSystemMessage

      public void addSystemMessage(String content)
      Adds a system message to the conversation history.
      Parameters:
      content - the message content

    • addMessage

      public void addMessage(Map<String,Object> message)
      Adds a structured message to the history.
      Parameters:
      message - the message map

    • getHistory

      public List<Map<String,Object>> getHistory()
      Returns the full conversation history.
      Returns:
      unmodifiable list of messages

    • getRecentHistory

      public List<Map<String,Object>> getRecentHistory(int limit)
      Returns recent messages from the history.
      Parameters:
      limit - maximum messages to return
      Returns:
      list of recent messages

    • clearHistory

      public void clearHistory()
      Clears all conversation history.

    • prune

      public void prune()
      Prunes history to stay within token limit.

    • captureCheckpoint

      public MemoryManager.AgentCheckpoint captureCheckpoint(String label)
      Captures the current agent state as a checkpoint.
      Parameters:
      label - descriptive label for the checkpoint
      Returns:
      the created checkpoint

    • captureCheckpoint

      public MemoryManager.AgentCheckpoint captureCheckpoint()
      Captures checkpoint with auto-generated label.
      Returns:
      the created checkpoint

    • restoreCheckpoint

      public boolean restoreCheckpoint(MemoryManager.AgentCheckpoint checkpoint)
      Restores agent state from a checkpoint.
      Parameters:
      checkpoint - the checkpoint to restore
      Returns:
      true if restoration was successful

    • rollbackToLast

      public boolean rollbackToLast()
      Restores to the most recent checkpoint.
      Returns:
      true if a checkpoint was restored

    • rollbackToLabel

      public boolean rollbackToLabel(String label)
      Restores to a checkpoint by label.
      Parameters:
      label - the checkpoint label
      Returns:
      true if found and restored

    • getLastCheckpoint

      public Optional<MemoryManager.AgentCheckpoint> getLastCheckpoint()
      Returns the most recent checkpoint if available.
      Returns:
      optional containing the last checkpoint

    • getCheckpointHistory

      public List<MemoryManager.AgentCheckpoint> getCheckpointHistory()
      Returns all stored checkpoints.
      Returns:
      unmodifiable list of checkpoints (oldest first)

    • getCheckpointCount

      public int getCheckpointCount()
      Returns the number of stored checkpoints.
      Returns:
      checkpoint count

    • clearCheckpoints

      public void clearCheckpoints()
      Clears all stored checkpoints.

    • search

      public List<String> search(String query, int limit)
      Performs semantic search on the memory. Useful for RAG (Retrieval-Augmented Generation).
      Parameters:
      query - the search query
      limit - maximum results
      Returns:
      list of relevant memory segments

    • getAgentId

      public String getAgentId()
      Returns the agent ID.
      Returns:
      the agent ID

    • getMemoryStore

      public MemoryStore getMemoryStore()
      Returns the underlying memory store.
      Returns:
      the memory store

    • getMaxContextTokens

      public Integer getMaxContextTokens()
      Returns the configured max context tokens.
      Returns:
      max tokens or null if not set