MCP Memory Service (Rust Implementation)

MCP Memory Service（Rust 实现）

这是使用官方 Rust MCP SDK 实现的 Model Context Protocol（MCP）Memory Service。

功能

• 实现一个 MCP 服务，提供记忆存储和检索功能
• 通过 stdio 通信，便于与 MCP 客户端集成
• 提供以下工具：
  • store_memory：存储包含内容、标签和元数据的新记忆
  • retrieve_memory：检索与查询在语义上相似的记忆
  • search_by_tag：按标签搜索记忆
  • delete_memory：按哈希值删除记忆
• 支持多种存储后端：
  • 内存存储（用于测试和开发）
  • ChromaDB 存储（用于生产环境）
• 支持多种嵌入模型：
  • 虚拟嵌入生成器（用于测试和开发）
  • ONNX 嵌入模型（使用 ONNX Runtime 的基于 Transformer 的嵌入）

先决条件

• Rust 和 Cargo（1.75.0 或更高版本）
• Node.js 和 npm（用于通过 MCP inspector 进行测试）
• 可选：ChromaDB 服务（用于生产环境）

构建

cargo build

构建发布版本：

cargo build --release

运行

直接运行服务：

cargo run

使用环境变量进行配置：

MCP_MEMORY_STORAGE_BACKEND=chromadb \
MCP_MEMORY_CHROMA_PATH=/path/to/chroma \
MCP_MEMORY_EMBEDDING_MODEL=onnx \
cargo run

配置

可以使用环境变量配置服务：

变量	描述	默认值
MCP_MEMORY_STORAGE_BACKEND	存储后端（inmemory 或 chromadb）	inmemory
MCP_MEMORY_CHROMA_PATH	ChromaDB 数据目录的路径	平台特定的数据目录
MCP_MEMORY_CHROMA_URL	ChromaDB 服务的 URL（可选）	无
MCP_MEMORY_CHROMA_COLLECTION	ChromaDB 集合名称	memory_collection
MCP_MEMORY_EMBEDDING_MODEL	嵌入模型（dummy 或 onnx）	dummy
MCP_MEMORY_EMBEDDING_MODEL_PATH	ONNX 模型文件的路径（可选）	无
MCP_MEMORY_EMBEDDING_SIZE	嵌入向量大小	384
MCP_MEMORY_LOG_LEVEL	日志级别	info

测试

此项目包含一个简单的 Node.js 测试脚本，用于测试服务。运行测试：

npm test

项目结构

• src/main.rs：主服务实现
• src/config.rs：配置处理
• src/models.rs：数据模型
• src/storage/：存储实现
  • mod.rs：存储 trait 和内存实现
  • chroma.rs：ChromaDB 存储实现
• src/embeddings.rs：嵌入模型实现
• src/utils.rs：实用函数
• Cargo.toml：Rust 项目配置
• package.json：用于测试工具的 Node.js 项目配置
• test-mcp-server.js：用于以编程方式测试服务的简单 Node.js 脚本

实现细节

该服务使用 Rust MCP SDK 中的 rmcp crate 实现。它使用：

• #[tool(tool_box)] 属性宏来定义工具
• ServerHandler trait 来实现 MCP 协议
• stdio() 传输方式供 CLI 使用
• 使用 Tokio 进行并发操作的异步 Rust
• Tracing 进行结构化日志记录

使用示例

使用 cargo run 启动服务后，可以向其 stdin 发送 JSON-RPC 请求。例如：

{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocol_version":"2024-11-05","capabilities":{"tools":true},"client_info":{"name":"test-client","version":"0.1.0"}}}
{"jsonrpc":"2.0","id":2,"method":"list_tools","params":{}}
{"jsonrpc":"2.0","id":3,"method":"call_tool","params":{"name":"store_memory","arguments":{"content":"This is a test memory","tags":["test","example"]}}}

ONNX 嵌入实现

ONNX 嵌入模型实现使用 ONNX Runtime 运行基于 Transformer 的模型，以生成嵌入。工作原理如下：

1. 模型加载：OnnxEmbeddingGenerator 从指定路径加载预训练的 ONNX 格式 Transformer 模型和分词器。

2. 分词：使用 HuggingFace tokenizers 库对输入文本进行分词，将文本转换为 token ID、注意力掩码和 token 类型 ID。

3. 推理：将分词后的输入传递给 ONNX 模型，由模型生成 Transformer 的隐藏状态。

4. 均值池化：使用均值池化（由注意力掩码加权）处理最后一个隐藏状态，以生成固定大小的嵌入向量。

5. 归一化：对生成的嵌入进行 L2 归一化，以确保向量幅值一致。

若要使用 ONNX 嵌入模型：

1. 使用 HuggingFace 的 transformers.onnx 等工具，将 Transformer 模型（如 BERT、RoBERTa 等）导出为 ONNX 格式。
2. 将分词器保存为与模型位于同一目录下的 tokenizer.json 文件，或指定单独的路径。
3. 设置环境变量：

   MCP_MEMORY_EMBEDDING_MODEL=onnx
   MCP_MEMORY_EMBEDDING_MODEL_PATH=/path/to/model.onnx
   MCP_MEMORY_EMBEDDING_SIZE=768  # Adjust based on your model's output size
   

在 MCP 客户端中注册

若要在 MCP 客户端（如 Claude）中注册此服务，需要将其添加到客户端的 MCP 配置中。以下是 JSON 配置示例：

{
  "mcpServers": {
    "memory-service": {
      "command": "/path/to/mcp-rust-server",
      "args": [],
      "env": {
        "MCP_MEMORY_STORAGE_BACKEND": "inmemory",
        "MCP_MEMORY_EMBEDDING_MODEL": "onnx",
        "MCP_MEMORY_EMBEDDING_MODEL_PATH": "/path/to/model.onnx",
        "MCP_MEMORY_EMBEDDING_SIZE": "768",
        "MCP_MEMORY_LOG_LEVEL": "info"
      },
      "disabled": false,
      "autoApprove": []
    }
  }
}

对于 Claude Desktop，此配置应添加到：

• macOS：~/Library/Application Support/Claude/claude_desktop_config.json
• Windows：%APPDATA%\Claude\claude_desktop_config.json
• Linux：~/.config/Claude/claude_desktop_config.json

对于 Claude VSCode 扩展，此配置应添加到：

• ~/Library/Application Support/Code/User/globalStorage/saoudrizwan.claude-dev/settings/cline_mcp_settings.json

未来改进

• 增加对更多 Transformer 架构的支持
• 增加更多存储后端
• 增加更多嵌入模型
• 增加更多内存管理工具
• 增加身份验证和授权
• 增加对更多传输方式（HTTP、WebSocket 等）的支持

架构

graph TD
    subgraph "External Interactions"
        direction LR
        Client([MCP Client]) -. JSON-RPC .-> StdioTransport[Stdio Transport]
        EnvVars[(Environment Variables)] -.-> Config
        ChromaDBServer[(ChromaDB Server)] <-. HTTP API .-> ChromaStorage
        OnnxModel[("ONNX Model File (.onnx)")] <-. Loads .-> OnnxEmbed
        TokenizerFile[("Tokenizer File (tokenizer.json)")] <-. Loads .-> OnnxEmbed
    end

    subgraph "MCP Memory Service (Rust Application)"
        direction TB
        StdioTransport -- Forwards Requests --> ServerCore{MCP Server Core} 

        ServerCore -- Reads --> Config(Configuration) 
        ServerCore -- Instantiates --> EmbeddingImpl{{Selected Embedding Generator}}
        ServerCore -- Instantiates --> StorageImpl{{Selected Storage Backend}}
        ServerCore -- Uses Tool Impls --> ToolLogic(Tool Logic: store, retrieve, search, delete)

        ToolLogic -- Uses --> StorageImpl
        ToolLogic -- Uses --> EmbeddingImpl
        ToolLogic -- Uses --> Models(Data Models) 
        ToolLogic -- Uses --> Utils(Utilities) 

        %% FIX: Added IDs EL and SL to subgraphs below
        subgraph EL [Embedding Layer]
            direction TB
            EmbeddingImpl -- Is an instance of --> EmbeddingTrait(EmbeddingGenerator Trait)
            DummyEmbed(DummyEmbeddingGenerator) -- Implements --> EmbeddingTrait
            OnnxEmbed(OnnxEmbeddingGenerator) -- Implements --> EmbeddingTrait
            OnnxEmbed -- Uses Lib --> OrtLib[ort crate]
            OnnxEmbed -- Uses Lib --> TokenizersLib[tokenizers crate]
        end

        subgraph SL [Storage Layer]
            direction TB
            StorageImpl -- Is an instance of --> StorageTrait(MemoryStorage Trait)
            InMemoryStorage(InMemoryStorage) -- Implements --> StorageTrait
            ChromaStorage(ChromaMemoryStorage) -- Implements --> StorageTrait
            StorageImpl -- Uses --> EmbeddingImpl
            StorageImpl -- Uses --> Models
            ChromaStorage -- Uses Lib --> ReqwestLib[reqwest crate]
        end

    end

    %% Styling
    classDef external fill:#f9f,stroke:#333,stroke-width:1px;
    classDef rust_app fill:#e6ffed,stroke:#333,stroke-width:1px;
    classDef trait fill:#lightblue,stroke:#333,stroke-width:1px;
    classDef impl fill:#lightgrey,stroke:#333,stroke-width:1px;
    classDef module fill:#whitesmoke,stroke:#333,stroke-width:1px;
    classDef lib fill:#cornsilk,stroke:#333,stroke-width:1px;


    class Client,EnvVars,ChromaDBServer,OnnxModel,TokenizerFile external;
    class StdioTransport,ServerCore,ToolLogic,Config,Models,Utils rust_app;
    class EmbeddingTrait,StorageTrait trait;
    class EmbeddingImpl,StorageImpl,DummyEmbed,OnnxEmbed,InMemoryStorage,ChromaStorage impl;
    class OrtLib,TokenizersLib,ReqwestLib lib;
    %% FIX: Used subgraph IDs EL and SL below
    class EL,SL module

许可证

本项目采用 MIT 许可证授权，详情请参阅 LICENSE 文件。