🧠 MCP - Titan Memory Server 实现

由 @jasonkneen 和 @ExpressionsBot 合作完成

在 X 上关注我们

• jasonkneen
• megaprompt

本实现受到 Google Research 的论文 "Generative AI for Programming: A Common Task Framework" 的启发。该服务器提供了一个神经记忆系统，能够学习和预测序列，并通过记忆向量保持状态，遵循了研究中提出的改进代码生成和理解的原则。

📚 研究背景

本实现借鉴了 Google Research 论文（Muennighoff 等人，2024）中提出的概念，该论文介绍了一种评估和改进代码生成模型的框架。Titan Memory 服务器实现了论文中的关键概念：

• 增强记忆的序列学习
• 用于新奇检测的惊喜度量
• 用于稳定学习的流形优化
• 通过记忆向量维持状态

这些特性与论文的目标一致，即通过更好的记忆和状态管理来改进代码理解和生成。

🚀 特性

• 可配置维度的神经记忆模型
• 序列学习和预测
• 惊喜度量计算
• 模型持久化（保存/加载）
• 记忆状态管理
• 完全的 MCP 工具集成

📦 安装

# Install dependencies
npm install

# Build the project
npm run build

# Run tests
npm test

🛠️ 可用的 MCP 工具

1. 🎯 init_model

使用自定义配置初始化 Titan Memory 模型。

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

2. 📚 train_step

使用当前和下一个状态向量执行单个训练步骤。

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

3. 🔄 forward_pass

使用输入向量通过模型运行前向传递。

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

4. 💾 save_model

将模型保存到指定路径。

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

5. 📂 load_model

从指定路径加载模型。

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

6. ℹ️ get_status

获取当前模型状态和配置。

{} // No parameters required

7. 🔄 train_sequence

对一系列向量进行模型训练。

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

🌟 使用示例

// 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, /* ... */]
});

🔧 技术细节

• 使用 TensorFlow.js 进行高效的张量操作
• 使用流形优化以实现稳定的训练
• 实现了用于新奇检测的惊喜度量
• 通过适当的张量清理进行内存管理
• 使用 TypeScript 实现类型安全
• 全面的错误处理

🧪 测试

项目包括全面的测试，涵盖：

• 模型初始化和配置
• 训练和前向传递操作
• 记忆状态管理
• 模型持久化
• 边界情况和错误处理
• 张量清理和内存管理

运行测试：

npm test

🔍 实现说明

• 所有张量操作都包裹在 tf.tidy() 中以确保正确的内存管理
• 实现了详细的错误处理并附带详细的错误信息
• 使用类型安全的 MCP 工具定义
• 在操作之间维护记忆状态
• 通过 epsilon 容差处理浮点精度问题

📝 许可证

MIT 许可证 - 随意使用和修改！