🧠 mcp - titan memory server implementation

colaboration between @jasonkneen and @expressionsbot

follow us on x

• jasonkneen
• megaprompt

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!