MCP 内存服务 (Rust 实现)

这是模型上下文协议 (MCP) 内存服务的 Rust 实现，使用了官方的 Rust MCP SDK。

特性

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

前提条件

• Rust 和 Cargo (1.75.0 或更高版本)
• Node.js 和 npm（用于使用 MCP 检查器进行测试）
• 可选：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（可选）	None
MCP_MEMORY_CHROMA_COLLECTION	ChromaDB 集合名称	memory_collection
MCP_MEMORY_EMBEDDING_MODEL	嵌入模型 (dummy 或 onnx)	dummy
MCP_MEMORY_EMBEDDING_MODEL_PATH	ONNX 模型文件的路径（可选）	None
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 启动服务器后，您可以将 JSON-RPC 请求发送到其 stdin。 例如：

{"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 模型和一个 tokenizer。

2. Tokenization: 输入文本使用 HuggingFace tokenizers 库进行 tokenization，该库将文本转换为 token ID、注意力掩码和 token 类型 ID。

3. 推理: tokenized 的输入传递给 ONNX 模型，该模型生成 Transformer 的隐藏状态。

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

5. 归一化: 对生成的嵌入进行 L2 归一化，以确保一致的向量大小。

要使用 ONNX 嵌入模型：

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

   MCP_MEMORY_EMBEDDING_MODEL=onnx
   MCP_MEMORY_EMBEDDING_MODEL_PATH=/path/to/model.onnx
   MCP_MEMORY_EMBEDDING_SIZE=768  # 根据模型的输出大小进行调整
   

向 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 文件。