Upload rag_pipeline.py

rag_pipeline.py

@@ -0,0 +1,312 @@
+"""
+RAG Pipeline for Internal Python Codebase
+
+Architecture based on:
+  - cAST (arxiv:2506.15655): AST-aware chunking via tree-sitter
+  - AllianceCoder (arxiv:2503.20589): API-first retrieval (signatures > similar code)
+  - CodeSage-v2 / Jina-Code-v2 for embeddings
+  - Gorilla (arxiv:2305.15334): retriever-aware generation pattern
+
+Components:
+  1. CodebaseIndexer — parses repo, chunks via AST, extracts API signatures
+  2. CodeRetriever — semantic search over code chunks and API signatures
+  3. ContextBuilder — assembles retrieval context for the LLM prompt
+
+Usage:
+  # Index and search a codebase
+  python rag_pipeline.py /path/to/repo --query "authentication token validation"
+
+  # Save/load index for fast startup
+  python rag_pipeline.py /path/to/repo --save-index ./index
+  python rag_pipeline.py /path/to/repo --load-index ./index --query "user permissions"
+"""
+
+import os
+import json
+import hashlib
+from pathlib import Path
+from dataclasses import dataclass, field
+from typing import Optional
+import numpy as np
+
+
+@dataclass
+class CodeChunk:
+    """A semantically meaningful piece of code."""
+    content: str
+    file_path: str
+    start_line: int
+    end_line: int
+    chunk_type: str  # "function", "class", "method", "module_level"
+    name: Optional[str] = None
+    parent_class: Optional[str] = None
+    signature: Optional[str] = None
+    docstring: Optional[str] = None
+    imports: list = field(default_factory=list)
+
+    @property
+    def id(self) -> str:
+        return hashlib.md5(f"{self.file_path}:{self.start_line}:{self.end_line}".encode()).hexdigest()
+
+    @property
+    def metadata_str(self) -> str:
+        parts = []
+        if self.chunk_type in ("function", "method"):
+            parts.append(f"Function {self.name}")
+            if self.signature: parts.append(f"with signature {self.signature}")
+            if self.docstring: parts.append(f"described as: {self.docstring}")
+            if self.parent_class: parts.append(f"in class {self.parent_class}")
+        elif self.chunk_type == "class":
+            parts.append(f"Class {self.name}")
+            if self.docstring: parts.append(f"described as: {self.docstring}")
+        parts.append(f"in file {self.file_path}")
+        return " ".join(parts)
+
+
+class ASTChunker:
+    """Parse Python files using AST and extract semantically meaningful chunks."""
+
+    def __init__(self, max_chunk_chars: int = 3000):
+        self.max_chunk_chars = max_chunk_chars
+
+    def chunk_file(self, file_path: str, source_code: str) -> list[CodeChunk]:
+        import ast
+        chunks = []
+        try:
+            tree = ast.parse(source_code)
+        except SyntaxError:
+            return [CodeChunk(content=source_code, file_path=file_path,
+                            start_line=1, end_line=source_code.count("\\n") + 1,
+                            chunk_type="module_level", name=Path(file_path).stem)]
+
+        lines = source_code.splitlines()
+        module_imports = []
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.Import, ast.ImportFrom)):
+                module_imports.append(ast.get_source_segment(source_code, node) or "")
+
+        for node in ast.iter_child_nodes(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                chunks.append(self._extract_function(node, source_code, lines, file_path, module_imports))
+            elif isinstance(node, ast.ClassDef):
+                chunks.append(self._extract_class(node, source_code, lines, file_path, module_imports))
+                for item in node.body:
+                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                        chunks.append(self._extract_function(item, source_code, lines, file_path, module_imports, parent_class=node.name))
+
+        module_lines = []
+        top_level_defs = {n.lineno for n in ast.iter_child_nodes(tree)
+                          if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))}
+        for i, line in enumerate(lines, 1):
+            if i not in top_level_defs:
+                in_def = False
+                for node in ast.iter_child_nodes(tree):
+                    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                        if hasattr(node, 'end_lineno') and node.lineno <= i <= node.end_lineno:
+                            in_def = True; break
+                if not in_def:
+                    module_lines.append(line)
+
+        if module_lines:
+            module_content = "\\n".join(module_lines).strip()
+            if module_content:
+                chunks.append(CodeChunk(content=module_content, file_path=file_path,
+                    start_line=1, end_line=len(lines), chunk_type="module_level",
+                    name=Path(file_path).stem, imports=module_imports))
+        return chunks
+
+    def _extract_function(self, node, source, lines, file_path, module_imports, parent_class=None):
+        import ast
+        start, end = node.lineno, node.end_lineno or node.lineno
+        content = "\\n".join(lines[start - 1:end])
+        args = []
+        for arg in node.args.args:
+            arg_str = arg.arg
+            if arg.annotation:
+                ann = ast.get_source_segment(source, arg.annotation)
+                if ann: arg_str += f": {ann}"
+            args.append(arg_str)
+        sig = f"def {node.name}({', '.join(args)})"
+        if node.returns:
+            ret = ast.get_source_segment(source, node.returns)
+            if ret: sig += f" -> {ret}"
+        return CodeChunk(content=content, file_path=file_path, start_line=start, end_line=end,
+            chunk_type="method" if parent_class else "function", name=node.name,
+            parent_class=parent_class, signature=sig, docstring=(ast.get_docstring(node) or "")[:500],
+            imports=module_imports)
+
+    def _extract_class(self, node, source, lines, file_path, module_imports):
+        import ast
+        start, end = node.lineno, node.end_lineno or node.lineno
+        bases = [ast.get_source_segment(source, b) or "" for b in node.bases]
+        sig = f"class {node.name}" + (f"({', '.join(bases)})" if bases else "")
+        full_content = "\\n".join(lines[start - 1:end])
+        return CodeChunk(content=full_content, file_path=file_path, start_line=start, end_line=end,
+            chunk_type="class", name=node.name, signature=sig,
+            docstring=(ast.get_docstring(node) or "")[:500], imports=module_imports)
+
+
+class CodeRetriever:
+    """Semantic search over code chunks using sentence-transformers embeddings."""
+
+    def __init__(self, embedding_model: str = "jinaai/jina-embeddings-v2-base-code"):
+        self.embedding_model_name = embedding_model
+        self.model = None
+        self.chunks: list[CodeChunk] = []
+        self.embeddings: Optional[np.ndarray] = None
+        self.signature_embeddings: Optional[np.ndarray] = None
+
+    def load_model(self):
+        if self.model is None:
+            try:
+                from sentence_transformers import SentenceTransformer
+                self.model = SentenceTransformer(self.embedding_model_name, trust_remote_code=True)
+            except ImportError:
+                self.model = "tfidf"
+
+    def index_chunks(self, chunks: list[CodeChunk]):
+        self.load_model()
+        self.chunks = chunks
+        if self.model == "tfidf":
+            from sklearn.feature_extraction.text import TfidfVectorizer
+            self.tfidf = TfidfVectorizer(max_features=10000, ngram_range=(1, 2))
+            self.tfidf_matrix = self.tfidf.fit_transform([c.content + " " + c.metadata_str for c in chunks])
+            return
+        contents = [c.content for c in chunks]
+        self.embeddings = self.model.encode(contents, batch_size=32, show_progress_bar=True, normalize_embeddings=True)
+        metadata = [c.metadata_str for c in chunks]
+        self.signature_embeddings = self.model.encode(metadata, batch_size=32, show_progress_bar=True, normalize_embeddings=True)
+
+    def search(self, query: str, top_k: int = 5, search_type: str = "hybrid") -> list[tuple[CodeChunk, float]]:
+        self.load_model()
+        if self.model == "tfidf":
+            query_vec = self.tfidf.transform([query])
+            scores = (self.tfidf_matrix @ query_vec.T).toarray().flatten()
+            top_indices = scores.argsort()[-top_k:][::-1]
+            return [(self.chunks[i], float(scores[i])) for i in top_indices if scores[i] > 0]
+        query_emb = self.model.encode([query], normalize_embeddings=True)
+        if search_type == "code":
+            scores = (query_emb @ self.embeddings.T).flatten()
+        elif search_type == "semantic":
+            scores = (query_emb @ self.signature_embeddings.T).flatten()
+        else:
+            scores = 0.4 * (query_emb @ self.embeddings.T).flatten() + 0.6 * (query_emb @ self.signature_embeddings.T).flatten()
+        top_indices = scores.argsort()[-top_k:][::-1]
+        return [(self.chunks[i], float(scores[i])) for i in top_indices]
+
+    def save_index(self, path: str):
+        os.makedirs(path, exist_ok=True)
+        if self.embeddings is not None:
+            np.save(os.path.join(path, "embeddings.npy"), self.embeddings)
+            np.save(os.path.join(path, "signature_embeddings.npy"), self.signature_embeddings)
+        with open(os.path.join(path, "chunks.json"), "w") as f:
+            json.dump([{"content": c.content, "file_path": c.file_path, "start_line": c.start_line,
+                "end_line": c.end_line, "chunk_type": c.chunk_type, "name": c.name,
+                "parent_class": c.parent_class, "signature": c.signature,
+                "docstring": c.docstring, "imports": c.imports} for c in self.chunks], f)
+
+    def load_index(self, path: str):
+        self.embeddings = np.load(os.path.join(path, "embeddings.npy"))
+        self.signature_embeddings = np.load(os.path.join(path, "signature_embeddings.npy"))
+        with open(os.path.join(path, "chunks.json")) as f:
+            self.chunks = [CodeChunk(**d) for d in json.load(f)]
+
+
+class CodebaseIndexer:
+    """Index an entire Python codebase."""
+
+    def __init__(self, repo_path: str, embedding_model: str = "jinaai/jina-embeddings-v2-base-code",
+                 max_chunk_chars: int = 3000, exclude_patterns: list[str] = None):
+        self.repo_path = Path(repo_path)
+        self.chunker = ASTChunker(max_chunk_chars=max_chunk_chars)
+        self.retriever = CodeRetriever(embedding_model=embedding_model)
+        self.exclude_patterns = exclude_patterns or ["__pycache__", ".git", ".venv", "venv", "node_modules"]
+
+    def index(self) -> CodeRetriever:
+        py_files = sorted(f for f in self.repo_path.rglob("*.py")
+                         if not any(e in f.parts for e in self.exclude_patterns) and f.stat().st_size < 100_000)
+        print(f"Found {len(py_files)} Python files")
+        all_chunks = []
+        for fpath in py_files:
+            try:
+                source = fpath.read_text(encoding="utf-8", errors="ignore")
+                chunks = self.chunker.chunk_file(str(fpath.relative_to(self.repo_path)), source)
+                all_chunks.extend(chunks)
+            except Exception as e:
+                print(f"  Warning: {fpath}: {e}")
+        print(f"Extracted {len(all_chunks)} chunks")
+        self.retriever.index_chunks(all_chunks)
+        return self.retriever
+
+
+class ContextBuilder:
+    """Build retrieval context for LLM prompts (AllianceCoder pattern)."""
+
+    def __init__(self, retriever: CodeRetriever, max_context_tokens: int = 4000):
+        self.retriever = retriever
+        self.max_context_chars = max_context_tokens * 4
+
+    def build_context(self, query: str, current_file_content: Optional[str] = None,
+                      current_file: Optional[str] = None, top_k: int = 5) -> str:
+        context_parts = []
+        total_chars = 0
+        if current_file_content:
+            in_ctx = self._extract_in_context_deps(current_file_content)
+            if in_ctx:
+                context_parts.append(f"# In-context dependencies from {current_file or 'current file'}:\\n{in_ctx}")
+                total_chars += len(in_ctx)
+        for chunk, score in self.retriever.search(query, top_k=top_k, search_type="hybrid"):
+            if total_chars >= self.max_context_chars: break
+            if chunk.signature:
+                entry = f"# From {chunk.file_path} (relevance: {score:.2f})\\n{chunk.signature}\\n"
+                if chunk.docstring: entry += f'    \"\"\"{chunk.docstring[:200]}\"\"\"\\n'
+            else:
+                entry = f"# From {chunk.file_path}:{chunk.start_line}-{chunk.end_line}\\n{chunk.content[:1000]}\\n"
+            context_parts.append(entry)
+            total_chars += len(entry)
+        return "\\n\\n".join(context_parts)
+
+    def _extract_in_context_deps(self, source: str) -> str:
+        import ast
+        try: tree = ast.parse(source)
+        except SyntaxError: return ""
+        deps = []
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    deps.append(f"import {alias.name}" + (f" as {alias.asname}" if alias.asname else ""))
+            elif isinstance(node, ast.ImportFrom):
+                deps.append(f"from {node.module} import {', '.join(a.name for a in node.names)}")
+        for node in ast.iter_child_nodes(tree):
+            if isinstance(node, ast.ClassDef):
+                deps.append(f"class {node.name}: ...")
+            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                deps.append(f"def {node.name}({', '.join(a.arg for a in node.args.args)}): ...")
+        return "\\n".join(deps)
+
+    def format_prompt_with_context(self, user_query: str, context: str, system_prompt: Optional[str] = None) -> list[dict]:
+        if not system_prompt:
+            system_prompt = "You are an expert Python programmer with access to our internal codebase via retrieval search."
+        user_content = f"{user_query}\\n\\n--- Retrieved context ---\\n{context}\\n--- End ---" if context else user_query
+        return [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_content}]
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("repo_path")
+    parser.add_argument("--query", "-q", default=None)
+    parser.add_argument("--model", default="jinaai/jina-embeddings-v2-base-code")
+    parser.add_argument("--save-index", default=None)
+    parser.add_argument("--load-index", default=None)
+    args = parser.parse_args()
+
+    if args.load_index:
+        retriever = CodeRetriever(args.model)
+        retriever.load_index(args.load_index)
+    else:
+        retriever = CodebaseIndexer(args.repo_path, embedding_model=args.model).index()
+    if args.save_index: retriever.save_index(args.save_index)
+    if args.query:
+        for i, (chunk, score) in enumerate(retriever.search(args.query, top_k=5)):
+            print(f"[{score:.3f}] {chunk.file_path}/{chunk.name} ({chunk.chunk_type})")