🧠 Task Memory Engine (TME)

Task Memory Engine (TME) is a structured memory framework for LLM-based agents, enabling multi-step task planning, rollback, replacement, and graph-based reasoning.

📄 About This Repository

This repository contains prototype code for two research versions of TME:

• v1: Tree + graph memory framework with slot-based task tracking
  ↪︎ Paper: Task Memory Engine (TME): A Structured Memory Framework with Graph-Aware Extensions for Multi-Step LLM Agent Tasks

• v2: Spatial memory system with rollback, replacement, DAG dependencies, and memory-aware QA
  ↪︎ Paper: Task Memory Engine: Spatial Memory for Robust Multi-Step LLM Agents

⚠️ Disclaimer: This is a reference implementation aligned with the above papers. v1 and v2 are conceptually related but structurally distinct. The repository is under active development, and modules may change before final release.

🚀 Quick Start

1. Install Dependencies

pip install openai
pip install python-dotenv  # if using .env to manage openai keys (recommended)

2. Set Your OPENAI API Key

export OPENAI_API_KEY=your_key_here

Or use a .env file:

OPENAI_API_KEY=your_key_here

3. Run Example Cases

Test cases are .json files in the cases/ directory, each containing a sequence of user instructions.
Test cases are running on ChatGPT-4o model.

Run Commands:

# Run with default classifier (general)
python run_case.py cases/trip_planning_case.json
python run_case.py cases/cooking_case.json
python run_case.py cases/meeting_scheduling_case.json

# Run cart case with specialized intent_classifier
python run_case.py cases/cart_editing_case.json --mode cart

🧠 Key Features

• Task Memory Tree (TMT): Hierarchical, structured task memory
• Rollback / Replace: Update or revert previous decisions
• Graph Reasoning (DAG): Non-linear dependencies between subtasks
• Instruction Decomposer: LLM-based substep splitting
• TRIM (Task Relation Inference Module): Classify task relations (merge, depend, rollback, etc.)
• Memory-Aware QA: Answer queries like “what’s currently in memory?”

🏗️ System Architecture

TME processes user inputs into a structured graph of subtasks, preserving history, dependencies, and intent transitions. Below are the architectural diagrams for v1 and v2:

• TME v1 Architecture: Illustrates the tree + graph memory framework with slot-based task tracking.
  

• TME v2 Architecture: Depicts the structured memory system with DAG dependencies and memory-aware QA.
  
  

Directory Structure

TME-Agent/
├── run_case.py                         # Main script to execute test cases
├── cases/                              # JSON input files for test scenarios
├── assets/                             # Architecture diagrams
├── v2/
│   ├── TaskMemoryStructure.py          # TaskNode & TaskMemoryTree logic
│   ├── input_splitter.py               # LLM-based instruction decomposition
│   ├── trim.py                         # Task relation reasoning (TRIM)
│   ├── intent_classifier_general.py    # General classifier (default)
│   └── intent_classifier_specific/
│       └── intent_classifier_cart.py   # Cart-specific classifier
├── citation.bib
└── README.md

🛡️ License & Usage

This project is licensed under the Polyform Noncommercial License 1.0.0.

• Free for academic and personal use.
• For commercial use, please contact the author directly for a license. 📧 Contact: biubiutomato@gmail.com

🌟 Star, Cite, Collaborate

If this project inspires or assists you, please consider:

• ⭐ Starring the repository
• 🧵 Opening discussions or issues
• 📚 Citing the relevant paper(s)

Let’s build memory-aware LLM agents together!