工具检索增强服务

MCP RAG是一个轻量级的AI SDK包装器,通过基于图的向量搜索智能索引和检索MCP工具,显著减少大型工具库的上下文开销。

  • by rocket-connect
  • TypeScript
  • 本地部署
  • 6 星标
  • 更新于
在 GitHub 查看 rocket-connect/mcp-rag

使用说明

MCP抹布

一个围绕AI SDK的轻量级包装器,使用基于图的矢量搜索智能地索引和检索MCP工具。

它做什么

MCP RAG将您的MCP工具集索引到图结构中,并使用Neo4j驱动的向量搜索从大型集合中检索相关的工具子集。这大大减少了使用广泛的工具库时的上下文开销。

基准测试

与基线工具选择相比,MCP RAG在效率和性能方面都有所提高,同时保持了相同的精度水平。

基准方法: 测试模拟了一个具有5个连续提示的真实对话,随着上下文的积累,每个提示都会触发一个不同的工具,反映了现实世界中的多回合交互。所有测试都使用GitHub MCP Server的完整工具集(90多种工具)来表示真实的大规模工具选择场景。

看看布丁里的证据🍰:

基础工具选择结果 -基线方法将所有工具传递给模型。

RAG工具选择结果 -基于RAG的矢量搜索智能滤波。

查看测试套件 -完成基准测试和测试用例。

安装

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

设置您的OpenAI API密钥:

export OPENAI_API_KEY=your_key_here

快速开始

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?

这 sync() 该方法执行工具与Neo4j的完全同步,创建语义搜索所需的图结构。下面是引擎盖下发生的事情:

  1. 创建矢量索引:使用1536维嵌入(OpenAI的 text-embedding-3-small 型号)

  2. 生成嵌入:对于工具集中的每个工具,它都会为以下对象创建嵌入:

    • 工具本身(名称+描述)
    • 每个参数(名称+描述)
    • 返回类型
  3. 构建图形结构:在Neo4j中创建具有以下关系的图:

    • ToolSet 将工具组合在一起的节点
    • Tool 节点及其嵌入
    • Parameter 通过以下方式连接到工具的节点 HAS_PARAM 关系
    • ReturnType 通过以下方式连接到工具的节点 RETURNS 关系
  4. 设计理念:同步过程使用 MERGE 操作,因此多次运行它不会创建重复项。如果工具集已更改,它将更新现有节点。

何时调用它:

  • 初始客户端创建后(首次使用前需要)
  • 在添加或删除工具后 addTool() 或 removeTool()
  • 强制重新索引您的工具

同步过程经过优化,仅在必要时运行-后续调用 generateText() 除非您明确调用,否则不会重新同步 sync() 再次或修改工具集。

What does rag.generateText() do?

这 generateText() 方法是围绕AI SDK的智能包装器 generateText 添加自动工具选择的功能。工作流程如下:

  1. 确保迁移:如果工具尚未编入索引,则自动调用同步过程

  2. 语义工具选择:

    • 为您的提示生成嵌入
    • 执行Neo4j向量相似性搜索,以找到最相关的工具
    • 默认情况下,最多可选择10个工具(可通过以下方式配置 maxActiveTools)
    • 您可以通过传递来覆盖此内容 activeTools 显式数组
  3. 调用AI SDK:仅将选定的工具子集传递给AI SDK的原生工具 generateText 功能以及您的提示和任何其他选项

  4. 返回完整结果:返回完整的AI SDK结果,封装在 GenerateTextResultWrapper 对象,使您可以访问:

    • 模型进行的工具调用
    • 令牌使用统计
    • 响应内容
    • 所有其他AI SDK元数据

主要优势:

  • 减小上下文大小:仅将相关工具发送到LLM,以节省令牌
  • 更好的性能:更少的工具意味着更快的响应时间
  • 相同的AI SDK体验:接受所有标准AI SDK参数并返回熟悉的结果结构

例子

GitHub MCP服务器演示

想看看MCP RAG的行动吗?查看我们的完整示例,该示例演示了如何使用GitHub MCP Server的93个工具进行智能工具选择:

📖 查看GitHub示例→

此示例显示:

  • 如何模拟和索引所有93个GitHub MCP服务器工具
  • 矢量相似性搜索——选择前10个最相关的工具
  • 具有详细调试输出的真实工具选择
  • 具有不同提示的交互式测试

非常适合了解MCP RAG如何在大型工具集中减少上下文开销!

特性

  • 基于图的索引 –工具与其关系和元数据一起编入索引
  • 向量搜索 –Neo4j支持的工具检索语义搜索
  • AI SDK兼容 –与您现有的AI SDK设置配合使用的插入式包装器
  • 选择性加载 –只加载每个请求所需的工具

API 参考

createMCPRag(config)

创建MCP RAG客户端。

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>,
  }
})

客户端方法

方法 说明
sync(options?) 将工具同步到Neo4j。退货 { hash: string }
generateText(options) 使用语义工具选择生成文本
getActiveTools(options) 获取语义上选定的工具,而不生成文本
addTool(name, tool) 添加工具(更新哈希值,需要重新同步)
removeTool(name) 删除工具(更新哈希,需要重新同步)
getTools() 获取所有已注册的工具
getToolsetHash() 获取当前工具集哈希值
getToolsetByHash(hash) 从Neo4j检索工具集信息
deleteToolsetByHash(hash) 从Neo4j中删除工具集

getActiveTools(options)

获取语义上选定的工具,用于您自己的AI SDK调用:

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?',
})

工具集哈希

哈希唯一地标识了用于更改检测和多版本支持的工具集版本。

Hashes如何工作

  1. 序列化 –工具被深度克隆(不包括 execute 功能)
  2. 排序 –工具按名称排序,所有嵌套键递归排序
  3. 哈希 –传递给哈希函数的JSON字符串
// Hash changes when tools change
const hash1 = rag.getToolsetHash()
rag.addTool('newTool', myTool)
const hash2 = rag.getToolsetHash() // Different from hash1

默认哈希

使用位哈希返回 toolset-<hex>:

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

自定义哈希函数

对于浏览器环境或自定义要求:

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('')
  },
})

哈希属性

  • 确定性的 –相同的工具集总是产生相同的哈希值
  • 订单独立 –工具/属性顺序不影响哈希
  • 变化敏感 –任何定义更改都会产生不同的哈希值

迁移

迁移将工具定义同步到Neo4j,为向量搜索创建图形结构。

基本迁移

await rag.sync() // Migrates if needed

迁移流动

  1. 检查 –确定是否需要迁移(Neo4j中是否存在哈希?)
  2. 索引 –创建向量索引(1536维,余弦相似度)
  3. 嵌入 –为工具、参数、返回类型生成嵌入
  4. 商店 在Neo4j中创建图形结构

自定义迁移挂钩

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
    },
  },
})

多版本工具集

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

工具集生命周期

// 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)

许可证

麻省理工学院 rconnect.tech