MCP 代码搜索与代码上下文工具：AI Coding Agent 如何理解整个代码库

import ast
import os
from dataclasses import dataclass, field
from typing import List, Dict, Optional
from pathlib import Path

@dataclass
class CodeChunk:
    """代码块：AST 解析后的语义单元"""
    filepath: str
    chunk_type: str  # 'function', 'class', 'module'
    name: str
    source_code: str
    line_start: int
    line_end: int
    docstring: Optional[str] = None
    imports: List[str] = field(default_factory=list)
    # 用于检索的特征
    keywords: List[str] = field(default_factory=list)

class ASTCodeSplitter:
    """基于 AST 的 Python 代码分块器"""
    
    def __init__(self):
        self.chunks: List[CodeChunk] = []
    
    def parse_file(self, filepath: str) -> List[CodeChunk]:
        """解析单个 Python 文件，提取函数和类"""
        with open(filepath, 'r', encoding='utf-8') as f:
            source = f.read()
        
        try:
            tree = ast.parse(source, filename=filepath)
        except SyntaxError:
            return []
        
        source_lines = source.split('\n')
        file_chunks = []
        
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                chunk = self._extract_function(node, filepath, source_lines)
                file_chunks.append(chunk)
            elif isinstance(node, ast.ClassDef):
                chunk = self._extract_class(node, filepath, source_lines)
                file_chunks.append(chunk)
        
        # 提取文件级 import
        imports = [n.names[0].name for n in ast.walk(tree)
                   if isinstance(n, ast.Import)]
        for chunk in file_chunks:
            chunk.imports = imports
        
        self.chunks.extend(file_chunks)
        return file_chunks
    
    def _extract_function(self, node: ast.FunctionDef,
                          filepath: str, lines: List[str]) -> CodeChunk:
        start = node.lineno - 1
        end = node.end_lineno or (start + 1)
        docstring = ast.get_docstring(node)
        
        return CodeChunk(
            filepath=filepath,
            chunk_type="function",
            name=node.name,
            source_code='\n'.join(lines[start:end]),
            line_start=node.lineno,
            line_end=node.end_lineno or node.lineno,
            docstring=docstring,
            keywords=self._extract_keywords(node, docstring)
        )
    
    def _extract_class(self, node: ast.ClassDef,
                       filepath: str, lines: List[str]) -> CodeChunk:
        start = node.lineno - 1
        end = node.end_lineno or (start + 1)
        docstring = ast.get_docstring(node)
        
        return CodeChunk(
            filepath=filepath,
            chunk_type="class",
            name=node.name,
            source_code='\n'.join(lines[start:end]),
            line_start=node.lineno,
            line_end=node.end_lineno or node.lineno,
            docstring=docstring,
            keywords=self._extract_keywords(node, docstring)
        )
    
    def _extract_keywords(self, node, docstring: Optional[str]) -> List[str]:
        """从函数/类名和文档字符串提取关键词"""
        keywords = []
        if docstring:
            keywords.extend(docstring.lower().split())
        keywords.append(node.name.lower())
        return list(set(keywords))

    def index_directory(self, dir_path: str, pattern: str = "*.py") -> List[CodeChunk]:
        """索引整个目录"""
        for py_file in Path(dir_path).rglob(pattern):
            self.parse_file(str(py_file))
        return self.chunks

# 使用示例
if __name__ == "__main__":
    splitter = ASTCodeSplitter()
    chunks = splitter.index_directory("./src")
    print(f"索引完成: 共 {len(chunks)} 个代码块")
    for c in chunks[:5]:
        print(f"  [{c.chunk_type}] {c.name} in {c.filepath}")

import math
from collections import Counter
from typing import List, Tuple
from dataclasses import dataclass

@dataclass
class SearchResult:
    chunk_name: str
    filepath: str
    score: float
    snippet: str

class CodeSearchEngine:
    """基于 TF-IDF 的代码搜索引擎"""
    
    def __init__(self):
        self.documents: List[str] = []  # 每个代码块的文本
        self.chunk_map: dict = {}  # 索引 -> CodeChunk
        self.idf: dict = {}  # 逆文档频率
        self.tfidf_matrix: List[dict] = []  # 每个文档的 TF-IDF 向量
    
    def build_index(self, chunks):
        """从代码块构建搜索索引"""
        for i, chunk in enumerate(chunks):
            # 合并代码 + 文档串 + 关键词作为搜索文本
            text = f"{chunk.source_code} {chunk.docstring or ''} {' '.join(chunk.keywords)}"
            self.documents.append(text)
            self.chunk_map[i] = chunk
        
        # 计算 IDF
        self._compute_idf()
        
        # 计算 TF-IDF 矩阵
        self._compute_tfidf_matrix()
    
    def _tokenize(self, text: str) -> List[str]:
        """简单的代码分词：按非字母数字字符分割"""
        import re
        # 保留 camelCase 和 snake_case 的语义
        tokens = re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', text.lower())
        return tokens
    
    def _compute_idf(self):
        n_docs = len(self.documents)
        doc_freq = Counter()
        
        for doc in self.documents:
            tokens = set(self._tokenize(doc))
            for token in tokens:
                doc_freq[token] += 1
        
        for token, freq in doc_freq.items():
            self.idf[token] = math.log((n_docs + 1) / (freq + 1)) + 1
    
    def _compute_tfidf_matrix(self):
        for doc in self.documents:
            tokens = self._tokenize(doc)
            tf = Counter(tokens)
            total = len(tokens) or 1
            
            tfidf = {}
            for term, count in tf.items():
                tf_val = count / total
                idf_val = self.idf.get(term, 1.0)
                tfidf[term] = tf_val * idf_val
            
            self.tfidf_matrix.append(tfidf)
    
    def search(self, query: str, top_k: int = 10) -> List[SearchResult]:
        """语义搜索代码"""
        query_tokens = self._tokenize(query)
        query_tfidf = Counter()
        
        for token in query_tokens:
            idf_val = self.idf.get(token, 1.0)
            query_tfidf[token] = idf_val  # 假设查询中每个词出现一次
        
        # 计算每个文档与查询的余弦相似度
        scores = []
        for i, doc_tfidf in enumerate(self.tfidf_matrix):
            score = self._cosine_similarity(query_tfidf, doc_tfidf)
            if score > 0:
                chunk = self.chunk_map[i]
                scores.append(SearchResult(
                    chunk_name=chunk.name,
                    filepath=chunk.filepath,
                    score=score,
                    snippet=chunk.source_code[:200]
                ))
        
        scores.sort(key=lambda x: x.score, reverse=True)
        return scores[:top_k]
    
    def _cosine_similarity(self, vec1: Counter, vec2: Counter) -> float:
        """计算两个 TF-IDF 向量的余弦相似度"""
        all_terms = set(vec1.keys()) | set(vec2.keys())
        
        dot = sum(vec1.get(t, 0) * vec2.get(t, 0) for t in all_terms)
        norm1 = math.sqrt(sum(v**2 for v in vec1.values()))
        norm2 = math.sqrt(sum(v**2 for v in vec2.values()))
        
        if norm1 == 0 or norm2 == 0:
            return 0.0
        return dot / (norm1 * norm2)

# 使用示例
if __name__ == "__main__":
    # 假设已从 ASTCodeSplitter 获得 chunks
    # chunks = splitter.index_directory("./src")
    
    # engine = CodeSearchEngine()
    # engine.build_index(chunks)
    
    # 模拟搜索结果
    results = [
        SearchResult("authenticate_user", "src/auth.py", 0.85, "def authenticate_user(token):..."),
        SearchResult("UserSession", "src/models.py", 0.72, "class UserSession:..."),
        SearchResult("login_handler", "src/handlers.py", 0.68, "async def login_handler(request):..."),
    ]
    
    print("搜索结果：")
    for r in results:
        print(f"  [{r.score:.2f}] {r.chunk_name} in {r.filepath}")
        print(f"    {r.snippet}")
        print()

from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel
from typing import List, Optional
import json

# 创建 MCP Server
mcp = FastMCP(
    "code-search",
    version="1.0.0",
    description="MCP Server for semantic code search"
)

class CodeSearchResult(BaseModel):
    name: str
    filepath: str
    score: float
    snippet: str
    chunk_type: str

@mcp.tool()
async def code_search(
    query: str,
    repo_path: str = ".",
    top_k: int = 10,
    chunk_types: Optional[List[str]] = None
) -> str:
    """
    在代码库中执行语义搜索。
    
    Args:
        query: 自然语言搜索查询（如 "处理用户认证的函数"）
        repo_path: 代码库路径（默认当前目录）
        top_k: 返回结果数量（默认 10）
        chunk_types: 过滤代码块类型（可选，如 ["function", "class"]）
    
    Returns:
        JSON 格式的搜索结果，包含代码片段、文件路径和相关度分数
    """
    # 1. AST 分块
    from ast_parser import ASTCodeSplitter
    splitter = ASTCodeSplitter()
    chunks = splitter.index_directory(repo_path)
    
    # 2. 类型过滤
    if chunk_types:
        chunks = [c for c in chunks if c.chunk_type in chunk_types]
    
    # 3. 构建索引并搜索
    from search_engine import CodeSearchEngine
    engine = CodeSearchEngine()
    engine.build_index(chunks)
    
    results = engine.search(query, top_k=top_k)
    
    # 4. 格式化为 JSON
    output = [
        {
            "name": r.chunk_name if hasattr(r, 'chunk_name') else r.name,
            "filepath": r.filepath,
            "score": round(r.score, 3),
            "snippet": r.snippet[:500],
            "chunk_type": r.chunk_type if hasattr(r, 'chunk_type') else "unknown"
        }
        for r in results
    ]
    
    return json.dumps(output, ensure_ascii=False, indent=2)

@mcp.tool()
async def code_nav(
    repo_path: str = ".",
    max_depth: int = 3
) -> str:
    """
    浏览代码库的目录结构。
    
    Args:
        repo_path: 代码库路径
        max_depth: 最大目录深度
    
    Returns:
        代码库的树状结构
    """
    from pathlib import Path
    
    def build_tree(path: Path, depth: int, prefix: str = "") -> str:
        if depth > max_depth:
            return ""
        
        lines = []
        entries = sorted(path.iterdir(), key=lambda x: (not x.is_dir(), x.name))
        
        for i, entry in enumerate(entries):
            if entry.name.startswith(('.git', '__pycache__', '.venv')):
                continue
            
            is_last = (i == len(entries) - 1)
            connector = "└── " if is_last else "├── "
            lines.append(f"{prefix}{connector}{entry.name}")
            
            if entry.is_dir():
                extension = "    " if is_last else "│   "
            else:
                extension = ""
            
            if entry.is_dir() and depth < max_depth:
                lines.append(build_tree(entry, depth + 1, prefix + extension))
        
        return "\n".join(filter(None, lines))
    
    tree = build_tree(Path(repo_path), 0)
    return f"📁 {repo_path}\n{tree}"

@mcp.tool()
async def code_read(filepath: str, line_start: int = 1, line_end: int = 100) -> str:
    """
    读取指定文件的代码内容。
    
    Args:
        filepath: 文件路径
        line_start: 起始行号（1-based）
        line_end: 结束行号（1-based）
    
    Returns:
        代码内容，带行号
    """
    with open(filepath, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    selected = lines[line_start - 1:line_end]
    numbered = [f"{i + line_start:4d} | {line.rstrip()}" 
                for i, line in enumerate(selected)]
    
    return f"📄 {filepath} (行 {line_start}-{line_end})\n" + "\n".join(numbered)

# 启动 MCP Server
if __name__ == "__main__":
    mcp.run()