🧠 MCP - Titan Memory Server implementation

Colaboration between @jasonkneen and @ExpressionsBot

An implementation inspired by Google Research's paper "Generative AI for Programming: A Common Task Framework". This server provides a neural memory system that can learn and predict sequences while maintaining state through a memory vector, following principles outlined in the research for improved code generation and understanding.

📚 Research Background

This implementation draws from the concepts presented in the Google Research paper (Muennighoff et al., 2024) which introduces a framework for evaluating and improving code generation models. The Titan Memory Server implements key concepts from the paper:

Memory-augmented sequence learning
Surprise metric for novelty detection
Manifold optimization for stable learning
State maintenance through memory vectors

These features align with the paper's goals of improving code understanding and generation through better memory and state management.

🚀 Features

Neural memory model with configurable dimensions
Sequence learning and prediction
Surprise metric calculation
Model persistence (save/load)
Memory state management
Full MCP tool integration

📦 Installation

# Install dependencies
npm install

# Build the project
npm run build

# Run tests
npm test

🛠️ Available MCP Tools

1. 🎯 init_model

Initialize the Titan Memory model with custom configuration.

{
  inputDim?: number;  // Input dimension (default: 64)
  outputDim?: number; // Output/Memory dimension (default: 64)
}

2. 📚 train_step

Perform a single training step with current and next state vectors.

{
  x_t: number[];    // Current state vector
  x_next: number[]; // Next state vector
}

3. 🔄 forward_pass

Run a forward pass through the model with an input vector.

{
  x: number[]; // Input vector
}

4. 💾 save_model

Save the model to a specified path.

{
  path: string; // Path to save the model
}

5. 📂 load_model

Load the model from a specified path.

{
  path: string; // Path to load the model from
}

6. ℹ️ get_status

Get current model status and configuration.

{} // No parameters required

7. 🔄 train_sequence

Train the model on a sequence of vectors.

{
  sequence: number[][]; // Array of vectors to train on
}

🌟 Example Usage

// Initialize model
await callTool('init_model', { inputDim: 64, outputDim: 64 });

// Train on a sequence
const sequence = [
  [1, 0, 0, /* ... */],
  [0, 1, 0, /* ... */],
  [0, 0, 1, /* ... */]
];
await callTool('train_sequence', { sequence });

// Run forward pass
const result = await callTool('forward_pass', {
  x: [1, 0, 0, /* ... */]
});

🔧 Technical Details

Built with TensorFlow.js for efficient tensor operations
Uses manifold optimization for stable learning
Implements surprise metric for novelty detection
Memory management with proper tensor cleanup
Type-safe implementation with TypeScript
Comprehensive error handling

🧪 Testing

The project includes comprehensive tests covering:

Model initialization and configuration
Training and forward pass operations
Memory state management
Model persistence
Edge cases and error handling
Tensor cleanup and memory management

Run tests with:

npm test

🔍 Implementation Notes

All tensor operations are wrapped in tf.tidy() for proper memory management
Implements proper error handling with detailed error messages
Uses type-safe MCP tool definitions
Maintains memory state between operations
Handles floating-point precision issues with epsilon tolerance

📝 License

MIT License - feel free to use and modify as needed!