rag_tool.py

from __future__ import annotations

import os
from pathlib import Path
import ast
import chromadb

CHROMA_DB_DIR = ".chroma_db"
COLLECTION_NAME = "codebase_rag"
TEXT_EXTENSIONS = {".py", ".md", ".txt", ".js", ".ts", ".json"}
IGNORE_PARTS = {".git", "__pycache__", "node_modules", "venv", "env", ".venv", ".chroma_db", "dist", "build", ".egg-info"}

_client = chromadb.PersistentClient(path=CHROMA_DB_DIR)
_collection = None


def _get_collection():
    global _collection
    if _collection is None:
        _collection = _client.get_or_create_collection(name=COLLECTION_NAME)
    return _collection


def _should_ignore(path: Path) -> bool:
    return any(part in IGNORE_PARTS for part in path.parts)


def _iter_text_files(workspace_dir: str):
    for root, _, files in os.walk(workspace_dir):
        root_path = Path(root)
        if _should_ignore(root_path):
            continue
        for name in files:
            file_path = root_path / name
            if _should_ignore(file_path):
                continue
            if file_path.suffix.lower() in TEXT_EXTENSIONS:
                yield file_path


def build_index_for_workspace(workspace_dir: str = "."):
    """Scan the workspace, chunk text files, and store them in ChromaDB."""
    workspace_root = Path(workspace_dir).resolve()
    print(f"🔍 正在为目录 [{workspace_root}] 构建代码向量索引...")

    global _collection
    try:
        _client.delete_collection(name=COLLECTION_NAME)
    except Exception:
        pass
    _collection = _client.get_or_create_collection(name=COLLECTION_NAME)

    documents: list[str] = []
    metadatas: list[dict] = []
    ids: list[str] = []
    chunk_id = 0

    # for file_path in _iter_text_files(str(workspace_root)):
    #     try:
    #         content = file_path.read_text(encoding="utf-8", errors="replace")
    #     except Exception as e:
    #         print(f"⚠️ 跳过文件 {file_path}: {e}")
    #         continue

    #     lines = content.split("\n")
    #     chunk_size = 80
    #     for start in range(0, len(lines), chunk_size):
    #         chunk_text = "\n".join(lines[start:start + chunk_size]).strip()
    #         if not chunk_text:
    #             continue
    #         documents.append(chunk_text)
    #         metadatas.append({"file_path": str(file_path), "start_line": start + 1})
    #         ids.append(f"chunk_{chunk_id}")
    #         chunk_id += 1
    for file_path in _iter_text_files(str(workspace_root)):
        try:
            content = file_path.read_text(encoding="utf-8", errors="replace")
        except Exception as e:
            print(f"⚠️ 跳过文件 {file_path}: {e}")
            continue

        # ==========================================
        # 🚀 [V2.0 升级] AST 语义切块 (Semantic Chunking)
        # ==========================================
        if file_path.suffix.lower() == ".py":
            # 针对 Python 文件，使用 AST 精准提取函数和类
            try:
                tree = ast.parse(content)
                extracted_chunks = False
                
                for node in tree.body:
                    # 我们只提取【函数】、【异步函数】和【类】作为独立的 Chunk
                    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                        # ast.get_source_segment 可以完美切出这个节点对应的完整代码 (需 Python 3.8+)
                        chunk_text = ast.get_source_segment(content, node)
                        if chunk_text and chunk_text.strip():
                            documents.append(chunk_text.strip())
                            metadatas.append({"file_path": str(file_path), "start_line": node.lineno})
                            ids.append(f"chunk_{chunk_id}")
                            chunk_id += 1
                            extracted_chunks = True
                
                # 如果这个 py 文件里全是全局变量或脚本逻辑，没有函数/类，我们再退化为把整个文件存进去
                if not extracted_chunks and content.strip():
                    documents.append(content.strip()[:4000]) # 限制下长度防止超载
                    metadatas.append({"file_path": str(file_path), "start_line": 1})
                    ids.append(f"chunk_{chunk_id}")
                    chunk_id += 1

            except SyntaxError:
                # 容错机制：如果代码本身有语法错误导致 AST 崩溃，就退化为把整个文件存进去
                if content.strip():
                    documents.append(content.strip()[:4000])
                    metadatas.append({"file_path": str(file_path), "start_line": 1})
                    ids.append(f"chunk_fallback_{chunk_id}")
                    chunk_id += 1
        else:
            # 针对 Markdown 或其他非 Python 文件，依然可以使用按行切块或整存策略
            lines = content.split("\n")
            chunk_size = 100
            for start in range(0, len(lines), chunk_size):
                chunk_text = "\n".join(lines[start:start + chunk_size]).strip()
                if not chunk_text:
                    continue
                documents.append(chunk_text)
                metadatas.append({"file_path": str(file_path), "start_line": start + 1})
                ids.append(f"chunk_txt_{chunk_id}")
                chunk_id += 1
        # ==========================================
    if not documents:
        print("⚠️ 未找到任何需要索引的代码文件。")
        return

    print(f"⏳ 正在计算向量并存入数据库，共 {len(documents)} 个代码片段...")
    _collection.upsert(documents=documents, metadatas=metadatas, ids=ids)
    print("✅ 向量索引构建完成！")


def search_code(query: str, n_results: int = 3) -> str:
    """Semantic search over the local codebase."""
    print(f"\n🕵️‍♂️ [RAG 检索] 正在搜索: '{query}'")
    collection = _get_collection()
    if collection.count() == 0:
        return "数据库为空，请先调用 build_index_for_workspace 构建索引。"

    results = collection.query(query_texts=[query], n_results=n_results)
    documents = results.get("documents", [[]])
    metadatas = results.get("metadatas", [[]])
    if not documents or not documents[0]:
        return "未找到相关代码。"

    response_text = f"基于查询 '{query}'，找到了以下相关代码片段：\n\n"
    for doc, meta in zip(documents[0], metadatas[0]):
        response_text += f"--- 来源文件: {meta['file_path']} (第 {meta['start_line']} 行起) ---\n"
        response_text += f"{doc}\n\n"
    return response_text


if __name__ == "__main__":
    print("🚀 开始初始化工作区外脑...")
    build_index_for_workspace()
    print(search_code("AST 静态语法拦截器在哪？"))