feat(memory): implement core memory and agentic components

[keli-wen]

Overview

This pull request introduces a foundational memory system for our agentic "brain" component. This is the first step towards building stateful, multi-step agents by providing a structured way to record and replay an agent's history.

Refer to #64

classDiagram
    direction LR

    class AgentMemory {
        system_prompt: SystemPromptStep
        steps: list[MemoryStep]
        +reset()
        +get_succinct_steps()
        +get_full_steps()
        +replay(logger, detailed)
    }

    class MemoryStep {
        +dict()
        +to_messages(summary_mode)
    }

    class SystemPromptStep {
        system_prompt: str
    }

    class TaskStep {
        task: str
        task_images: list[Image]?
    }

    class PlanningStep {
        plan: str
        model_input_messages: list[ChatMessage]
        model_output_message: ChatMessage
        timing: Timing
    }

    class ActionStep {
        step_number: int
        model_input_messages: list[ChatMessage]?
        tool_calls: list[ToolCall]?
        model_output: str | list[dict]?
        observations: str?
        token_usage: TokenUsage?
    }

    class FinalAnswerStep {
        output: Any
    }

    class CallbackRegistry {
        +register(step_cls, callback)
        +callback(memory_step, **kwargs)
    }

    class ToolCall {
        name: str
        id: str
        arguments: Any
    }

    class ChatMessage {
        role: MessageRole
        content: str | list[dict]
    }

    class MessageRole
    class Timing
    class TokenUsage
    class AgentLogger

    AgentMemory --> SystemPromptStep
    AgentMemory --> MemoryStep
    AgentMemory --> AgentLogger
    AgentMemory --> CallbackRegistry

    MemoryStep <|-- SystemPromptStep
    MemoryStep <|-- TaskStep
    MemoryStep <|-- PlanningStep
    MemoryStep <|-- ActionStep
    MemoryStep <|-- FinalAnswerStep

    ActionStep --> ToolCall
    ActionStep --> ChatMessage
    ActionStep --> Timing
    ActionStep --> TokenUsage

    PlanningStep --> ChatMessage
    TaskStep --> ChatMessage

Loading

Key Changes

• Memory System (quantmind/brain/memory.py):

  • Introduced the Memory class to store a sequential timeline of an agent's actions, including tasks, tool calls, and planning steps.
  • Added a CallbackRegistry to allow for hooking into different stages of an agent's execution.

• Structured Data Models (quantmind/models/):

  • Created new dataclasses (ActionStep, TaskStep, PlanningStep, ChatMessage, etc.) to provide a strongly-typed representation of the agent's history and message formats.
  • These models handle serialization to dictionaries and conversion into message formats suitable for LLMs.

• Agentic Utilities (quantmind/utils/):

  • Added a suite of utilities for monitoring (AgentLogger, TokenUsage, Timing), error handling (AgentError subclasses), and data parsing.
  • The AgentLogger leverages the rich library for enhanced, readable console output.

• Examples & Tests:

  • Included a new example (examples/memory/basic_memory_usage.py) demonstrating how to construct and inspect a memory timeline.
  • Added comprehensive unit tests for the new memory and message models to ensure correctness.

• Documentation (AGENTS.md):

  • Updated the contributor guidelines to include a required comment style.

Implementation Notes

• The architecture separates the stateful Memory store from the immutable MemoryStep data models, promoting clarity and easier state management.
• The replay method in the Memory class, combined with AgentLogger, provides powerful debugging capabilities by visualizing the agent's entire run history.
• Utility functions in agentic_ext.py and messages.py are adapted from common agentic frameworks to provide a robust foundation for future development.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with $CATEGORY(xx): xxx (such as feat(tool): xxx, fix(source): xxx, docs(README): xxx)
• Related issue is referred in this PR
• The markdown and latex are rendered correctly.
• The code in PR is well-documented.
• The PR is complete and small, read the Google eng practice (CL equals to PR) to understand more about small PR.