Skip to content

rocket-connect/mcp-rag

Repository files navigation

MCP RAG

npm version CI License: MIT

A lightweight wrapper around AI SDK that intelligently indexes and retrieves MCP tools using graph-based vector search.

Neo4j Model

What It Does

MCP RAG indexes your MCP toolset into a graph structure and uses Neo4j-powered vector search to retrieve relevant tool subsets from large collections. This dramatically reduces context overhead when working with extensive tool libraries.

Benchmarks

MCP RAG sees improvements in both efficiency and performance compared to baseline tool selection, while maintaining the same level of accuracy.

Benchmark Methodology: Tests simulate a realistic conversation with 5 sequential prompts, each triggering a different tool as context accumulates—mirroring real-world multi-turn interactions. All tests use the complete toolset from the GitHub MCP Server (90+ tools) to represent authentic large-scale tool selection scenarios.

See the proof in the pudding 🍰:

Base Tool Selection Results - Baseline approach passing all tools to the model.

RAG Tool Selection Results - RAG-powered intelligent filtering with vector search.

View Test Suite - Complete benchmark implementation and test cases.

Installation

npm version

npm install @mcp-rag/client @ai-sdk/openai neo4j-driver ai

Set your OpenAI API key:

export OPENAI_API_KEY=your_key_here

Quick Start

import { createMCPRag } from '@mcp-rag/client'
import { openai } from '@ai-sdk/openai'
import neo4j from 'neo4j-driver'
import { tool } from 'ai'
import { z } from 'zod'

const driver = neo4j.driver(
  'neo4j://localhost:7687',
  neo4j.auth.basic('neo4j', 'password')
)

const rag = createMCPRag({
  model: openai('gpt-4o-mini'),
  openaiApiKey: process.env.OPENAI_API_KEY || '',
  neo4j: driver,
  tools: {
    searchDocs: tool({
      /* ... */
    }),
    queryDatabase: tool({
      /* ... */
    }),
    sendEmail: tool({
      /* ... */
    }),
    fetchWeather: tool({
      /* ... */
    }),
    analyzeImage: tool({
      /* ... */
    }),
    // ... hundreds more tools
  },
})

await rag.sync()
const result = await rag.generateText({
  prompt: 'Search for API docs',
})
What does rag.sync() do?

The sync() method performs a complete synchronization of your tools to Neo4j, creating the graph structure needed for semantic search. Here's what happens under the hood:

  1. Creates Vector Index: Sets up a Neo4j vector index for similarity search using 1536-dimensional embeddings (OpenAI's text-embedding-3-small model)

  2. Generates Embeddings: For each tool in your toolset, it creates embeddings for:

    • The tool itself (name + description)
    • Each parameter (name + description)
    • The return type
  3. Builds Graph Structure: Creates a graph in Neo4j with the following relationships:

    • ToolSet nodes that group tools together
    • Tool nodes with their embeddings
    • Parameter nodes connected to tools via HAS_PARAM relationships
    • ReturnType nodes connected to tools via RETURNS relationships
  4. Idempotent by Design: The sync process uses MERGE operations, so running it multiple times won't create duplicates. It will update existing nodes if the toolset has changed.

When to call it:

  • After initial client creation (required before first use)
  • After adding or removing tools with addTool() or removeTool()
  • To force a re-index of your tools

The sync process is optimized to only run when necessary - subsequent calls to generateText() won't re-sync unless you explicitly call sync() again or modify the toolset.

What does rag.generateText() do?

The generateText() method is a smart wrapper around the AI SDK's generateText function that adds automatic tool selection. Here's the workflow:

  1. Ensures Migration: Automatically calls the sync process if tools haven't been indexed yet

  2. Semantic Tool Selection:

    • Generates an embedding for your prompt
    • Performs a Neo4j vector similarity search to find the most relevant tools
    • By default, selects up to 10 tools (configurable via maxActiveTools)
    • You can override this by passing activeTools array explicitly
  3. Calls AI SDK: Passes only the selected subset of tools to the AI SDK's native generateText function along with your prompt and any additional options

  4. Returns Full Result: Returns the complete AI SDK result wrapped in a GenerateTextResultWrapper object, giving you access to:

    • Tool calls made by the model
    • Token usage statistics
    • Response content
    • All other AI SDK metadata

Key Benefits:

  • Reduced Context Size: Only relevant tools are sent to the LLM, saving tokens
  • Better Performance: Fewer tools mean faster response times
  • Same AI SDK Experience: Accepts all standard AI SDK parameters and returns familiar result structures
Tools Select Model

Examples

GitHub MCP Server Demo

Want to see MCP RAG in action? Check out our complete example that demonstrates intelligent tool selection with the GitHub MCP Server's 93 tools:

📖 View GitHub Example →

GitHub Tools in Neo4j Browser

This example shows:

  • How to mock and index all 93 GitHub MCP server tools
  • Vector similarity search selecting the top 10 most relevant tools
  • Real-world tool selection with detailed debug output
  • Interactive testing with different prompts

Perfect for understanding how MCP RAG reduces context overhead in large toolsets!

Features

  • Graph-based indexing – Tools are indexed with their relationships and metadata
  • Vector search – Neo4j-powered semantic search for tool retrieval
  • AI SDK compatible – Drop-in wrapper that works with your existing AI SDK setup
  • Selective loading – Only load the tools you need for each request

API Reference

createMCPRag(config)

Creates an MCP RAG client.

const rag = createMCPRag({
  model: LanguageModel,              // AI SDK model (required)
  neo4j: Driver,                     // Neo4j driver (required)
  tools: Record<string, Tool>,       // AI SDK tools (required)
  openaiApiKey: string,              // For embeddings (required)
  maxActiveTools?: number,           // Default: 10
  hashFunction?: (input: string) => string,  // Custom hash function
  dangerouslyAllowBrowser?: boolean, // Enable browser usage
  migration?: {
    shouldMigrate?: (session) => Promise<boolean>,
    migrate?: (session, tools) => Promise<void>,
    onBeforeMigrate?: (statements) => Promise<statements>,
  }
})

Client Methods

Method Description
sync(options?) Sync tools to Neo4j. Returns { hash: string }
generateText(options) Generate text with semantic tool selection
getActiveTools(options) Get semantically selected tools without generating text
addTool(name, tool) Add a tool (updates hash, requires re-sync)
removeTool(name) Remove a tool (updates hash, requires re-sync)
getTools() Get all registered tools
getToolsetHash() Get current toolset hash
getToolsetByHash(hash) Retrieve toolset info from Neo4j
deleteToolsetByHash(hash) Delete a toolset from Neo4j

getActiveTools(options)

Get semantically selected tools to use with your own AI SDK calls:

const { tools, names } = await rag.getActiveTools({
  prompt: 'What is the weather?',
  maxTools: 5, // optional, defaults to maxActiveTools
})

// Use with AI SDK directly
const result = await generateText({
  model: openai('gpt-4o'),
  tools,
  prompt: 'What is the weather?',
})

Toolset Hashes

Hashes uniquely identify toolset versions for change detection and multi-version support.

How Hashes Work

  1. Serialization – Tools are deep-cloned (excluding execute functions)
  2. Sorting – Tools sorted by name, all nested keys sorted recursively
  3. Hashing – JSON string passed to hash function
// Hash changes when tools change
const hash1 = rag.getToolsetHash()
rag.addTool('newTool', myTool)
const hash2 = rag.getToolsetHash() // Different from hash1

Default Hash

Uses a bitwise hash returning toolset-<hex>:

const hash = rag.getToolsetHash() // "toolset-a1b2c3d4"

Custom Hash Function

For browser environments or custom requirements:

const rag = createMCPRag({
  // ...
  hashFunction: input => {
    // Web Crypto API example
    const encoder = new TextEncoder()
    const data = encoder.encode(input)
    const hashBuffer = await crypto.subtle.digest('SHA-256', data)
    return Array.from(new Uint8Array(hashBuffer))
      .map(b => b.toString(16).padStart(2, '0'))
      .join('')
  },
})

Hash Properties

  • Deterministic – Same toolset always produces same hash
  • Order-independent – Tool/property order doesn't affect hash
  • Change-sensitive – Any definition change produces different hash

Migrations

Migration syncs tool definitions to Neo4j, creating the graph structure for vector search.

Basic Migration

await rag.sync() // Migrates if needed

Migration Flow

  1. Check – Determines if migration needed (hash exists in Neo4j?)
  2. Index – Creates vector index (1536 dimensions, cosine similarity)
  3. Embed – Generates embeddings for tools, parameters, return types
  4. Store – Creates graph structure in Neo4j

Custom Migration Hooks

const rag = createMCPRag({
  // ...
  migration: {
    // Override migration check
    shouldMigrate: async session => {
      const result = await session.run('...')
      return result.records.length === 0
    },

    // Custom migration logic
    migrate: async (session, tools) => {
      // Your migration code
    },

    // Intercept/modify migration statements
    onBeforeMigrate: async statements => {
      console.log('Migrating:', statements.length, 'statements')
      return statements // Return modified or original
    },
  },
})

Multi-Version Toolsets

Multiple toolset versions can coexist in Neo4j:

// Version 1
await rag.sync()
const v1Hash = rag.getToolsetHash()

// Version 2 (both exist in DB)
rag.addTool('newTool', myTool)
await rag.sync()
const v2Hash = rag.getToolsetHash()

// Manage versions
const v1Info = await rag.getToolsetByHash(v1Hash)
await rag.deleteToolsetByHash(v1Hash) // Clean up old version

Toolset Lifecycle

// 1. Setup
const { hash } = await rag.sync()

// 2. Retrieve
const info = await rag.getToolsetByHash(hash)
// { hash, updatedAt, toolCount, tools: [...] }

// 3. Update
rag.addTool('newTool', tool)
await rag.sync() // Creates new version

// 4. Cleanup
await rag.deleteToolsetByHash(oldHash)

License

MIT rconnect.tech

About

A lightweight wrapper around AI SDK that intelligently indexes and retrieves MCP tools using graph-based vector search.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Sponsor this project

 

Contributors