关键词

PyTorch文档搜索语义搜索代码示例开发工具

使用说明

项目简介

PyTorch文档搜索工具是一个基于Model Context Protocol (MCP) 的服务器实现，旨在为大型语言模型（LLM）客户端（如Claude Code）提供PyTorch官方文档的语义搜索能力。它允许用户通过自然语言查询快速找到相关的PyTorch API文档、代码示例、最佳实践和错误信息，从而提升开发效率。该工具利用向量嵌入和语义相似性技术，提供高质量的搜索结果，并支持多种传输协议（STDIO, SSE, UVX）以便灵活部署。

主要功能点

  • PyTorch文档语义搜索: 基于向量嵌入的语义搜索，能够理解用户查询的意图，返回更相关的文档片段。
  • 代码感知搜索结果: 区分代码块和文本内容，提供更精确的代码示例搜索结果。
  • MCP集成: 通过MCP协议与Claude Code等客户端无缝集成，作为工具扩展LLM的功能。
  • 多种传输协议支持: 支持STDIO、SSE等多种传输方式，适应不同的部署环境和客户端需求。
  • 可配置的搜索参数: 允许用户配置搜索结果数量和内容类型过滤（代码或文本）。

安装步骤

  1. 环境准备: 推荐使用conda创建独立的Python环境以隔离依赖。

    • 完整环境 (推荐): 
      conda env create -f environment.yml
conda activate pytorch_docs_search
    • 最小环境: 
      conda env create -f minimal_env.yml
conda activate pytorch_docs_search_min
    • 如果conda环境创建失败，请确保已安装conda，并检查'environment.yml'或'minimal_env.yml'文件是否存在且内容正确。

  2. API密钥配置: 本工具依赖OpenAI API生成文本嵌入向量，需要配置OpenAI API密钥。

    export OPENAI_API_KEY=你的OpenAI_API密钥
    • 请将'你的OpenAI_API密钥'替换为实际的OpenAI API密钥。如果未配置API密钥，服务器将无法启动。

服务器配置

本MCP服务器可以通过STDIO或SSE两种方式与MCP客户端连接。以下是针对不同连接方式的服务器配置信息，以JSON格式提供。MCP客户端需要根据选择的连接方式配置相应的 'command' 和 'args' 。

1. STDIO 连接 (本地开发常用)

{
  "server_name": "pytorch_docs_search_stdio",
  "command": "./run_mcp.sh",
  "args": []
}
  • 'server_name': 服务器名称，用于在MCP客户端中标识该服务器。可以自定义。
  • 'command': 服务器启动命令。'./run_mcp.sh' 脚本封装了使用STDIO传输方式启动服务器的命令。
  • 'args': 启动参数列表，对于STDIO传输方式，通常为空列表。

2. SSE 连接 (服务器部署常用)

{
  "server_name": "pytorch_docs_search_sse",
  "command": "python",
  "args": ["-m", "ptsearch.server", "--transport", "sse", "--host", "0.0.0.0", "--port", "5000"]
}
  • 'server_name': 服务器名称，用于在MCP客户端中标识该服务器。可以自定义。
  • 'command': 服务器启动命令，使用 'python -m ptsearch.server' 启动服务器模块。
  • 'args': 启动参数列表：
    • '"-m", "ptsearch.server"': 指定运行 'ptsearch.server' 模块。
    • '"–transport", "sse"': 指定使用SSE传输协议。
    • '"–host", "0.0.0.0"': 指定服务器监听地址为 '0.0.0.0'，允许来自任何IP的连接。可以根据需要修改为特定的IP地址或 'localhost' 。
    • '"–port", "5000"': 指定服务器监听端口为 '5000'。可以根据需要修改为其他可用端口。

3. UVX 连接 (打包分发)

{
  "server_name": "pytorch_docs_search_uvx",
  "command": "./run_mcp_uvx.sh",
  "args": []
}
  • 'server_name': 服务器名称，用于在MCP客户端中标识该服务器。可以自定义。
  • 'command': 服务器启动命令。'./run_mcp_uvx.sh' 脚本封装了使用UVX和SSE传输方式启动服务器的命令。
  • 'args': 启动参数列表，对于UVX传输方式，通常为空列表。

注意:

  • 选择哪种配置取决于MCP客户端支持的连接方式以及部署环境。通常本地开发可以使用STDIO，服务器部署可以使用SSE。
  • MCP客户端需要配置与以上JSON配置对应的服务器启动命令和参数，才能正确连接到MCP服务器。
  • 如果需要修改数据目录，可以在 'args' 中添加 '--data-dir <数据目录路径>' 参数，例如：'["-m", "ptsearch.server", "--transport", "sse", "--host", "0.0.0.0", "--port", "5000", "--data-dir", "./my_data"]'。

基本使用方法

  1. 启动MCP服务器: 根据选择的连接方式，使用相应的命令启动MCP服务器。例如，对于STDIO连接，运行 './run_mcp.sh' ；对于SSE连接，运行 'python -m ptsearch.server --transport sse --host 0.0.0.0 --port 5000' 。
  2. 在MCP客户端注册工具: 使用MCP客户端提供的命令或界面注册该工具。例如，对于Claude CLI，可以使用以下命令注册STDIO和SSE两种方式的工具：
    • STDIO: 
      claude mcp add search_pytorch_docs stdio ./run_mcp.sh
    • SSE: 
      claude mcp add search_pytorch_docs http://localhost:5000/events --transport sse
    • 请确保服务器已成功启动并监听在指定的地址和端口，再进行工具注册。
  3. 使用Claude Code搜索PyTorch文档: 在Claude Code等MCP客户端中，可以直接使用自然语言提问关于PyTorch的问题，例如： 
    How do I implement a custom dataset in PyTorch?
    Claude Code会自动调用PyTorch文档搜索工具，并返回相关的文档片段作为上下文信息。

直接命令行使用

除了作为MCP服务器，该工具也提供了直接的命令行搜索功能，方便用户在本地进行测试和使用。

# 使用STDIO传输方式启动服务器，并进入交互式搜索模式
python -m ptsearch.server --transport stdio --data-dir ./data

启动后，可以在命令行中输入搜索query进行PyTorch文档的搜索。

# 命令行单次查询示例
python -m ptsearch.server --transport stdio --data-dir ./data < "如何使用PyTorch定义一个卷积层？"

信息

分类

开发者工具

访问项目