""" 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})")