rag_pipeline.py

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