Python代码库分析RAG系统

该系统使用抽象语法树 (AST) 分析 Python 代码，将提取的信息（函数、类、调用、变量等）存储在 Weaviate 矢量数据库中，并通过模型上下文协议 (MCP) 服务器提供查询和理解代码库的工具。它利用 Google 的 Gemini 模型生成嵌入和自然语言描述/答案。

特征

• **代码扫描：**解析 Python 文件以识别代码元素（函数、类、导入、调用、赋值）及其关系。摘录：

  • 基本信息：名称、类型、文件路径、行号、代码片段、文档字符串。

  • 函数/方法详细信息：参数、返回类型、签名、装饰器。

  • 范围信息：父范围（类/函数）UUID、可读 ID（例如， file:type:name:line ）、基类名称。

  • 使用信息：范围内的属性访问、调用关系（部分跟踪）。

• **向量存储：**使用 Weaviate 存储代码元素及其向量嵌入（启用 LLM 生成时）。

• **LLM 扩展（可选且为后台）：**使用 Gemini 为函数和类生成语义描述和嵌入。此功能现在作为扫描后触发的后台任务运行，或手动运行。可通过.env文件启用/禁用。

• **自动细化（可选和后台）：**启用 LLM 生成时，使用上下文（调用者、被调用者、同级、相关变量）作为后台处理的一部分自动细化新函数/更新函数的描述。

• **RAG 问答：**使用检索增强生成回答有关代码库的自然语言问题（需要启用 LLM 功能并完成后台处理）。

• **用户说明：**允许用户向特定代码元素添加手动注释。

• **可视化：**根据存储的关系生成 MermaidJS 调用图。

• **MCP 服务器：**通过 MCP 工具公开分析和查询功能，管理代码库和活动代码库上下文。

• **文件观察器（集成）：**扫描代码库时自动启动（ scan_codebase ），选择另一个代码库时自动停止（ select_codebase ）或删除代码库时自动停止（ delete_codebase ）。当活动代码库的文件发生变更时，触发重新分析并更新数据库。也可以通过start_watcher和stop_watcher工具手动控制。

• **代码库依赖关系：**允许定义扫描代码库之间的依赖关系（ add_codebase_dependency ， remove_codebase_dependency ）。

• **跨代码库查询：**支持跨活动代码库及其声明的依赖项进行搜索（ find_element ）和提问（ ask_question ）。

Related MCP server: code2prompt-mcp

设置

1. **环境：**确保安装了 Python 3.10+ 和 Docker。

2. **Weaviate：**使用 Docker Compose 启动 Weaviate 实例：

   docker-compose up -d

3. **依赖项：**安装 Python 包：

   pip install -r requirements.txt

4. **API 密钥和配置：**在项目根目录下创建一个.env文件并添加您的 Gemini API 密钥。您还可以配置其他设置：

   # --- Required --- GEMINI_API_KEY=YOUR_API_KEY_HERE # --- Optional --- # Set to true to enable background LLM description generation and refinement GENERATE_LLM_DESCRIPTIONS=true # Max concurrent background LLM tasks (embeddings/descriptions/refinements) LLM_CONCURRENCY=5 # ANALYZE_ON_STARTUP is no longer used. Scanning is done via the scan_codebase tool. # Specify Weaviate connection details if not using defaults # WEAVIATE_HOST=localhost # WEAVIATE_PORT=8080 # WEAVIATE_GRPC_PORT=50051 # Specify alternative Gemini models if desired # GENERATION_MODEL_NAME="models/gemini-pro" # EMBEDDING_MODEL_NAME="models/embedding-001" # Adjust Weaviate batch size # WEAVIATE_BATCH_SIZE=100 # SEMANTIC_SEARCH_LIMIT=5 # SEMANTIC_SEARCH_DISTANCE=0.7 # Watcher polling interval (seconds) # WATCHER_POLLING_INTERVAL=5

5. **运行 MCP 服务器：**在单独的终端中启动服务器：

   python src/code_analysis_mcp/mcp_server.py

   （确保此终端保持运行以使工具可用）

架构概述

该系统分析 Python 代码，将提取的信息存储在 Weaviate 矢量数据库中，并通过模型上下文协议 (MCP) 服务器提供查询和理解代码库的工具。它利用 Google 的 Gemini 模型生成嵌入和自然语言描述/答案。

主要模块有：

• code_scanner.py ：查找 Python 文件，使用 AST 解析它们，提取结构元素（函数、类、导入、调用等），并为 Weaviate 准备数据。

• weaviate_client.py ：管理与 Weaviate 的连接，定义数据模式（ CodeFile 、 CodeElement 、 CodebaseRegistry ），并提供批量上传、查询、更新和删除数据的功能。

• rag.py ：实现检索增强生成 (RAG)，用于回答有关代码库的问题。它使用语义搜索查找相关的代码元素，并使用 LLM 来合成答案。

• mcp_server.py ：设置 FastMCP 服务器，管理CodebaseRegistry集合中的代码库，处理活动代码库上下文（ ACTIVE_CODEBASE_NAME ），集成文件监视逻辑（包括自动启动/停止），管理代码库依赖关系，并将分析功能作为 MCP 工具公开，并提供详细的参数描述。

• visualization.py ：根据存储的关系生成 MermaidJS 调用图。

系统使用 Weaviate 的多租户功能来管理CodeFile和CodeElement集合，其中租户 ID 是用户定义的codebase_name 。一个独立的非多租户CodebaseRegistry集合跟踪代码库元数据（名称、目录、状态、摘要、监视程序状态、依赖项）。服务器中的ACTIVE_CODEBASE_NAME全局变量确定查询的主要代码库租户。查询list_codebases （ find_element 、 ask_question ）可以选择搜索存储在注册表中的活动代码库及其声明的依赖项。list_codebases 工具可用于查看所有代码库的状态和依赖项。

后台 LLM 处理用于生成代码元素的语义描述和嵌入。这是一个可选功能，可以通过.env文件启用/禁用。

服务器运行后，可以使用标准 MCP 自检方法直接从 MCP 服务器检索有关可用工具及其参数的详细信息。