RAG工具检索系统

一个基于RAG的工具检索系统,用于MCP服务器,支持智能语义搜索和多服务器管理。

  • by fiestaxxl
  • Python
  • 本地部署
  • 1 星标
  • 更新于
在 GitHub 查看 fiestaxxl/rag_tools

使用说明

RAG工具

基于RAG的MCP(模型上下文协议)服务器工具检索系统。该模块支持对人工智能代理的MCP工具进行智能语义搜索。

特性

  • 语义搜索:使用自然语言查询查找相关工具
  • RAG管道:基于矢量的检索,可选择交叉编码器重新排序
  • 多服务器支持:管理来自多个MCP服务器的工具
  • 异步架构:完全异步实现,实现高性能
  • 评估框架:用于衡量检索质量的内置指标
  • 可扩展的:专为大型工具系列而设计

建筑

┌─────────────┐     ┌──────────────┐     ┌─────────────┐
│  MCP Server │────▶│   Parser     │────▶│  Text       │
│  (HTTP)     │     │              │     │  Builder    │
└─────────────┘     └──────────────┘     └──────┬──────┘
                                                │
                    ┌──────────────┐            ▼
                    │   Embedder   │◀────┌─────────────┐
                    │ (Local/API)  │     │  Indexer    │
                    └──────┬───────┘     └──────┬──────┘
                           │                    │
                           ▼                    ▼
                    ┌──────────────┐     ┌─────────────┐
                    │   Qdrant     │     │  PostgreSQL │
                    │   (Vectors)  │     │  (Registry) │
                    └──────────────┘     └─────────────┘

安装

pip install -r requirements.txt

要在另一个项目中安装该模块(例如,从Git存储库),您可以:

# Install from a local directory (editable mode)
pip install -e /path/to/rag_tools

# Install from a Git repository
pip install git+https://github.com/fiestaxxl/rag_tools.git

# Install with optional local models
pip install "rag_tools[local]"@git+https://github.com/fiestaxxl/rag_tools.git

快速开始

1.设置基础设施

cd docker
./setup.sh

2.初始化和使用

import asyncio
from rag_tools import create_manager

async def main():
    # Create manager
    manager = await create_manager()

    # Add an MCP server
    server = await manager.add_server(
        url="http://localhost:8080/mcp",
        name="my-server",
        description="My MCP server"
    )

    # Retrieve relevant tools
    results = await manager.retrieve_tools("search for academic papers")

    print(f"Found {len(results)} relevant tools:")
    for r in results:
        print(f"  - {r.name} (score: {r.score:.3f})")

    await manager.close()

asyncio.run(main())

用法示例

添加服务器

server = await manager.add_server(
    url="http://10.32.11.22:7331/mcp",
    name="openalex",
    description="OpenAlex academic search API",
    headers={"Authorization": "Bearer ..."},
    sync_tools=True  # Automatically fetch tools
)

检索工具

# Simple retrieval
results = await manager.retrieve_tools("search papers by author")

# With configuration
result = await manager.retrieve(
    query="download PDF documents",
    top_k=10,
    rerank=True,
    min_score=0.3
)

# Access detailed results
for r in result.results:
    print(f"{r.name}: {r.description[:100]}...")
    print(f"  Score: {r.score:.3f}, Rank: {r.rank}")

评估

from rag_tools import create_default_suite, RAGEvaluator

# Create evaluation suite
suite = create_default_suite()

# Or load from file
evaluator = RAGEvaluator(pipeline)
report = await evaluator.evaluate(suite)

# View results
evaluator.print_report(report)

配置

配置通过以下方式管理 config/settings.py 或环境变量,请检查 example/env.example:

# Environment variables
export QDRANT__URL=http://localhost:6333
export POSTGRES__HOST=localhost
export POSTGRES__PORT=5432
export POSTGRES__USER=rag_tools
export POSTGRES__PASSWORD=secret

关键设置

设置 默认值 说明
embedding.model_name 全MiniLM-L6-v2 嵌入式模型
reranker.model_name ms-marco-MiniLM-L-6-v2 Reranker模型
rag.default_top_k 10 默认检索限制
rag.rerank_top_k 5 重新排名后的结果

api参考

RAGTools管理器

工具管理和检索的主要类。

方法

  • add_server(url, name, ...) -添加MCP服务器
  • remove_server(server_id) -删除服务器
  • sync_server(server_id) -从服务器同步工具
  • add_tool(server_id, name, ...) -手动添加工具
  • remove_tool(tool_id) -移除工具
  • retrieve(query, ...) -检索完整结果
  • retrieve_tools(query, ...) -简化检索
  • evaluate(suite) -运行评估
  • get_stats() -获取统计数据

管道

检索管道支持:

  • 向量搜索 通过Qdrant
  • 交叉编码器重新排序 (可选)
  • 过滤 按服务器、标签
  • 分块 对于长篇描述

指标

评估模块提供:

  • Recall@K -检索到的相关工具的比例
  • Precision@K -检索到的相关工具的比例
  • 平均倒数排名 -平均倒数排名
  • 地图 -平均精度
  • NDCG@K -归一化贴现累计收益
  • 命中率 -是否有相关工具在前K

发展

运行测试

# Install dev dependencies
pip install pytest pytest-asyncio

# Run tests
pytest tests/

项目结构

rag_tools/
├── config/          # Configuration
├── storage/         # Database clients
├── ingestion/       # MCP ingestion
├── retrieval/       # RAG retrieval
├── evaluation/      # Metrics
├── tools/           # Tool management
└── main.py          # Main module

码头工人

服务定义见 docker/docker-compose.yml:

cd docker
docker-compose up -d

许可证

麻省理工学院