MCP Memory Service (Rust Implementation)

MCP Memory Service（Rust 実装）

これは、公式の Rust MCP SDK を使用した、Model Context Protocol（MCP）Memory Service の Rust 実装です。

機能

• メモリの保存と取得機能を提供する MCP サーバーを実装します
• MCP クライアントと簡単に連携できるよう、stdio 経由で通信します
• 次のツールを提供します：
  • store_memory：コンテンツ、タグ、メタデータを指定して新しいメモリを保存します
  • retrieve_memory：クエリと意味的に類似したメモリを取得します
  • search_by_tag：タグでメモリを検索します
  • delete_memory：ハッシュを指定してメモリを削除します
• 複数のストレージバックエンドをサポートします：
  • メモリ内ストレージ（テストおよび開発用）
  • ChromaDB ストレージ（本番環境用）
• 複数の埋め込みモデルをサポートします：
  • ダミー埋め込み生成器（テストおよび開発用）
  • ONNX 埋め込みモデル（ONNX Runtime を使用したトランスフォーマーベースの埋め込み）

前提条件

• 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（任意）	なし
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：ストレージトレイトとメモリ内実装
  • 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 クレートを使用して実装されています。以下を使用します：

• ツールを定義するための #[tool(tool_box)] 属性マクロ
• MCP プロトコルを実装するための ServerHandler トレイト
• CLI で使用するための stdio() トランスポート
• 並行処理のための 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 ライブラリを使用してトークン化され、トークン ID、アテンションマスク、トークンタイプ 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 ファイルを参照してください。