tarang 4.4.0__py3-none-any.whl

tarang/context/retriever.py

@@ -0,0 +1,270 @@
+"""
+Context Retriever - Unified interface for BM25 + KG retrieval.
+
+Combines BM25 keyword search with Symbol Graph expansion to
+provide rich, connected context for LLM queries.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from .bm25 import BM25Index, SearchResult
+from .chunker import Chunk
+from .graph import SymbolGraph, SymbolNode
+
+
+@dataclass
+class RetrievalResult:
+    """Result from context retrieval."""
+    chunks: List[Chunk]                         # Full code for direct matches
+    signatures: List[str]                       # Signatures for connected symbols
+    graph_context: Dict[str, Any]               # Relationship summary
+    stats: Dict[str, Any] = field(default_factory=dict)
+
+    def to_context_dict(self) -> Dict:
+        """Convert to dictionary for API payload."""
+        return {
+            "chunks": [
+                {
+                    "id": c.id,
+                    "file": c.file,
+                    "type": c.type,
+                    "name": c.name,
+                    "signature": c.signature,
+                    "content": c.content,
+                    "line_start": c.line_start,
+                    "line_end": c.line_end,
+                }
+                for c in self.chunks
+            ],
+            "signatures": self.signatures,
+            "graph": self.graph_context,
+        }
+
+    @property
+    def total_lines(self) -> int:
+        """Total lines of code in chunks."""
+        return sum(c.line_end - c.line_start + 1 for c in self.chunks)
+
+    @property
+    def is_empty(self) -> bool:
+        """Check if result has no content."""
+        return len(self.chunks) == 0
+
+
+class ContextRetriever:
+    """
+    Unified context retrieval using BM25 + Knowledge Graph.
+
+    Workflow:
+    1. BM25 search finds relevant chunks
+    2. KG expansion adds connected symbols (signatures only)
+    3. Returns combined context for LLM
+    """
+
+    def __init__(
+        self,
+        bm25_index: BM25Index,
+        symbol_graph: SymbolGraph,
+    ):
+        self.bm25 = bm25_index
+        self.graph = symbol_graph
+
+    @property
+    def is_ready(self) -> bool:
+        """Check if retriever has indexed data."""
+        return not self.bm25.is_empty
+
+    def retrieve(
+        self,
+        query: str,
+        hops: int = 1,
+        max_chunks: int = 10,
+        max_signatures: int = 20,
+    ) -> RetrievalResult:
+        """
+        Retrieve relevant context for a query.
+
+        Args:
+            query: User instruction or search query
+            hops: KG expansion hops (0=none, 1=direct, 2=2-level)
+            max_chunks: Maximum code chunks to return
+            max_signatures: Maximum connected signatures
+
+        Returns:
+            RetrievalResult with chunks, signatures, and graph context
+        """
+        # Step 1: BM25 search
+        search_results = self.bm25.search(query, k=max_chunks)
+
+        if not search_results:
+            return RetrievalResult(
+                chunks=[],
+                signatures=[],
+                graph_context={},
+                stats={"bm25_hits": 0, "expanded_symbols": 0},
+            )
+
+        # Extract chunks and symbol IDs
+        chunks = [r.chunk for r in search_results]
+        symbol_ids = [c.id for c in chunks]
+
+        # Step 2: KG expansion
+        signatures: List[str] = []
+        expanded_ids: set = set()
+
+        if hops > 0 and not self.graph.is_empty:
+            for sid in symbol_ids:
+                neighbors = self.graph.get_neighbors(sid, hops=hops)
+                for neighbor in neighbors:
+                    if neighbor.id not in symbol_ids and neighbor.id not in expanded_ids:
+                        expanded_ids.add(neighbor.id)
+                        signatures.append(neighbor.signature)
+
+                        if len(signatures) >= max_signatures:
+                            break
+
+                if len(signatures) >= max_signatures:
+                    break
+
+        # Step 3: Get graph context
+        all_ids = symbol_ids + list(expanded_ids)
+        graph_context = self.graph.get_graph_context(all_ids)
+
+        return RetrievalResult(
+            chunks=chunks,
+            signatures=signatures[:max_signatures],
+            graph_context=graph_context,
+            stats={
+                "bm25_hits": len(search_results),
+                "expanded_symbols": len(expanded_ids),
+                "total_chunks": len(chunks),
+                "total_signatures": len(signatures),
+            },
+        )
+
+    def retrieve_for_file(
+        self,
+        file_path: str,
+        hops: int = 1,
+    ) -> RetrievalResult:
+        """
+        Retrieve all context for a specific file.
+
+        Useful when user mentions a file explicitly.
+        """
+        chunks = self.bm25.get_chunks_for_file(file_path)
+
+        if not chunks:
+            return RetrievalResult(
+                chunks=[],
+                signatures=[],
+                graph_context={},
+            )
+
+        symbol_ids = [c.id for c in chunks]
+
+        # KG expansion
+        signatures: List[str] = []
+        expanded_ids: set = set()
+
+        if hops > 0 and not self.graph.is_empty:
+            for sid in symbol_ids:
+                neighbors = self.graph.get_neighbors(sid, hops=hops)
+                for neighbor in neighbors:
+                    if neighbor.id not in symbol_ids and neighbor.id not in expanded_ids:
+                        expanded_ids.add(neighbor.id)
+                        signatures.append(neighbor.signature)
+
+        graph_context = self.graph.get_graph_context(symbol_ids + list(expanded_ids))
+
+        return RetrievalResult(
+            chunks=chunks,
+            signatures=signatures,
+            graph_context=graph_context,
+        )
+
+    def retrieve_symbol(
+        self,
+        symbol_name: str,
+        hops: int = 1,
+    ) -> RetrievalResult:
+        """
+        Retrieve context for a specific symbol by name.
+
+        Searches for chunks matching the symbol name exactly.
+        """
+        # Search for the symbol
+        results = self.bm25.search(symbol_name, k=5)
+
+        # Filter to exact name matches
+        exact_matches = [
+            r for r in results
+            if r.chunk.name.lower() == symbol_name.lower()
+        ]
+
+        if not exact_matches:
+            # Fall back to partial matches
+            exact_matches = results[:3]
+
+        if not exact_matches:
+            return RetrievalResult(chunks=[], signatures=[], graph_context={})
+
+        chunks = [r.chunk for r in exact_matches]
+        symbol_ids = [c.id for c in chunks]
+
+        # KG expansion
+        signatures: List[str] = []
+        expanded_ids: set = set()
+
+        if hops > 0 and not self.graph.is_empty:
+            for sid in symbol_ids:
+                neighbors = self.graph.get_neighbors(sid, hops=hops)
+                for neighbor in neighbors:
+                    if neighbor.id not in symbol_ids and neighbor.id not in expanded_ids:
+                        expanded_ids.add(neighbor.id)
+                        signatures.append(neighbor.signature)
+
+        graph_context = self.graph.get_graph_context(symbol_ids + list(expanded_ids))
+
+        return RetrievalResult(
+            chunks=chunks,
+            signatures=signatures,
+            graph_context=graph_context,
+        )
+
+    def get_callers(self, symbol_id: str) -> List[SymbolNode]:
+        """Get all symbols that call this symbol."""
+        return self.graph.get_callers(symbol_id)
+
+    def get_callees(self, symbol_id: str) -> List[SymbolNode]:
+        """Get all symbols that this symbol calls."""
+        return self.graph.get_callees(symbol_id)
+
+
+def create_retriever(index_path: Path) -> Optional[ContextRetriever]:
+    """
+    Create a retriever from saved index files.
+
+    Args:
+        index_path: Path to .tarang/index/ directory
+
+    Returns:
+        ContextRetriever if index exists, None otherwise
+    """
+    bm25_path = index_path / "bm25.pkl"
+    graph_path = index_path / "graph.json"
+
+    bm25 = BM25Index()
+    graph = SymbolGraph()
+
+    # Load BM25 index
+    if not bm25.load(bm25_path):
+        return None
+
+    # Load graph (optional, retriever works without it)
+    graph.load(graph_path)
+
+    return ContextRetriever(bm25, graph)

tarang/context/skeleton.py

@@ -0,0 +1,282 @@
+"""
+Project Skeleton Generator - Lightweight project context for backend.
+
+Generates file tree and symbol information to send to the Orchestrator.
+"""
+from __future__ import annotations
+
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class SymbolDefinition:
+    """A symbol (function, class, method) in the project."""
+    name: str
+    kind: str  # function, class, method, variable
+    file: str
+    line: int
+    signature: Optional[str] = None
+
+
+@dataclass
+class ProjectSkeleton:
+    """Lightweight project map for context."""
+    file_tree: str
+    symbols: List[SymbolDefinition] = field(default_factory=list)
+    dependencies: Dict[str, List[str]] = field(default_factory=dict)
+    total_files: int = 0
+    total_lines: int = 0
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for API."""
+        return {
+            "file_tree": self.file_tree,
+            "symbols": [
+                {
+                    "name": s.name,
+                    "kind": s.kind,
+                    "file": s.file,
+                    "line": s.line,
+                    "signature": s.signature,
+                }
+                for s in self.symbols[:100]  # Limit symbols
+            ],
+            "dependencies": dict(list(self.dependencies.items())[:50]),
+            "total_files": self.total_files,
+            "total_lines": self.total_lines,
+        }
+
+
+class SkeletonGenerator:
+    """
+    Generate lightweight project skeleton for backend context.
+
+    Extracts file tree and symbol definitions without sending full code.
+    """
+
+    IGNORE_PATTERNS = [
+        "node_modules", ".git", "__pycache__", ".venv",
+        "venv", "dist", "build", ".next", "target",
+        ".tarang", ".pytest_cache", ".mypy_cache",
+        "*.pyc", "*.pyo", ".DS_Store",
+    ]
+
+    def __init__(self, project_root: Path):
+        self.project_root = project_root
+
+    def generate(self, max_depth: int = 4) -> ProjectSkeleton:
+        """
+        Generate project skeleton.
+
+        Args:
+            max_depth: Maximum directory depth for tree
+
+        Returns:
+            ProjectSkeleton with tree and symbols
+        """
+        file_tree = self._generate_tree(max_depth)
+        symbols = self._extract_symbols()
+        dependencies = self._analyze_dependencies()
+        total_files, total_lines = self._count_stats()
+
+        return ProjectSkeleton(
+            file_tree=file_tree,
+            symbols=symbols,
+            dependencies=dependencies,
+            total_files=total_files,
+            total_lines=total_lines,
+        )
+
+    def _should_ignore(self, path: Path) -> bool:
+        """Check if path should be ignored."""
+        name = path.name
+        for pattern in self.IGNORE_PATTERNS:
+            if pattern.startswith("*"):
+                if name.endswith(pattern[1:]):
+                    return True
+            elif pattern in str(path):
+                return True
+        return False
+
+    def _generate_tree(self, max_depth: int) -> str:
+        """Generate ASCII file tree."""
+        lines = [f"{self.project_root.name}/"]
+
+        def walk(path: Path, prefix: str = "", depth: int = 0):
+            if depth > max_depth:
+                return
+
+            try:
+                items = sorted(path.iterdir(), key=lambda x: (x.is_file(), x.name.lower()))
+            except PermissionError:
+                return
+
+            # Filter ignored items
+            items = [i for i in items if not self._should_ignore(i)]
+
+            for i, item in enumerate(items[:30]):  # Limit items per directory
+                is_last = i == len(items) - 1
+                connector = "└── " if is_last else "├── "
+
+                if item.is_dir():
+                    lines.append(f"{prefix}{connector}{item.name}/")
+                    extension = "    " if is_last else "│   "
+                    walk(item, prefix + extension, depth + 1)
+                else:
+                    lines.append(f"{prefix}{connector}{item.name}")
+
+        walk(self.project_root)
+        return "\n".join(lines[:200])  # Limit total lines
+
+    def _extract_symbols(self) -> List[SymbolDefinition]:
+        """Extract symbols using ctags if available."""
+        symbols = []
+
+        # Try ctags first
+        if self._has_ctags():
+            symbols = self._extract_with_ctags()
+            if symbols:
+                return symbols
+
+        # Fallback: Simple regex extraction for Python
+        symbols = self._extract_python_symbols()
+        return symbols
+
+    def _has_ctags(self) -> bool:
+        """Check if ctags is available."""
+        try:
+            subprocess.run(
+                ["ctags", "--version"],
+                capture_output=True,
+                timeout=5
+            )
+            return True
+        except (FileNotFoundError, subprocess.TimeoutExpired):
+            return False
+
+    def _extract_with_ctags(self) -> List[SymbolDefinition]:
+        """Extract symbols using universal-ctags."""
+        symbols = []
+
+        try:
+            result = subprocess.run(
+                [
+                    "ctags", "-R", "--output-format=json",
+                    "--languages=Python,JavaScript,TypeScript,Go,Rust",
+                    "--exclude=node_modules", "--exclude=.git",
+                    "--exclude=__pycache__", "--exclude=venv",
+                    "-f", "-", str(self.project_root)
+                ],
+                capture_output=True,
+                timeout=30
+            )
+
+            import json
+            for line in result.stdout.decode().strip().split("\n"):
+                if not line:
+                    continue
+                try:
+                    tag = json.loads(line)
+                    symbols.append(SymbolDefinition(
+                        name=tag.get("name", ""),
+                        kind=tag.get("kind", "unknown"),
+                        file=tag.get("path", ""),
+                        line=tag.get("line", 0),
+                        signature=tag.get("signature"),
+                    ))
+                except json.JSONDecodeError:
+                    continue
+
+        except (subprocess.TimeoutExpired, FileNotFoundError):
+            pass
+
+        return symbols[:500]
+
+    def _extract_python_symbols(self) -> List[SymbolDefinition]:
+        """Fallback: Extract Python symbols with regex."""
+        import re
+        symbols = []
+
+        func_pattern = re.compile(r'^(\s*)def\s+(\w+)\s*\(([^)]*)\)', re.MULTILINE)
+        class_pattern = re.compile(r'^class\s+(\w+)', re.MULTILINE)
+
+        for py_file in self.project_root.rglob("*.py"):
+            if self._should_ignore(py_file):
+                continue
+
+            try:
+                content = py_file.read_text(errors="replace")
+                rel_path = str(py_file.relative_to(self.project_root))
+
+                # Extract classes
+                for match in class_pattern.finditer(content):
+                    line_num = content[:match.start()].count("\n") + 1
+                    symbols.append(SymbolDefinition(
+                        name=match.group(1),
+                        kind="class",
+                        file=rel_path,
+                        line=line_num,
+                    ))
+
+                # Extract functions
+                for match in func_pattern.finditer(content):
+                    indent = match.group(1)
+                    name = match.group(2)
+                    args = match.group(3)
+                    line_num = content[:match.start()].count("\n") + 1
+
+                    kind = "method" if indent else "function"
+                    symbols.append(SymbolDefinition(
+                        name=name,
+                        kind=kind,
+                        file=rel_path,
+                        line=line_num,
+                        signature=f"({args})",
+                    ))
+
+            except (IOError, UnicodeDecodeError):
+                continue
+
+        return symbols[:500]
+
+    def _analyze_dependencies(self) -> Dict[str, List[str]]:
+        """Build import dependency graph for Python files."""
+        deps = {}
+
+        for py_file in self.project_root.rglob("*.py"):
+            if self._should_ignore(py_file):
+                continue
+
+            imports = []
+            try:
+                content = py_file.read_text(errors="replace")
+                for line in content.split("\n")[:100]:  # Only scan first 100 lines
+                    line = line.strip()
+                    if line.startswith("import ") or line.startswith("from "):
+                        imports.append(line)
+            except (IOError, UnicodeDecodeError):
+                continue
+
+            if imports:
+                rel_path = str(py_file.relative_to(self.project_root))
+                deps[rel_path] = imports[:20]
+
+        return deps
+
+    def _count_stats(self) -> tuple:
+        """Count total files and lines."""
+        total_files = 0
+        total_lines = 0
+
+        for f in self.project_root.rglob("*"):
+            if f.is_file() and not self._should_ignore(f):
+                total_files += 1
+                try:
+                    total_lines += len(f.read_text(errors="replace").split("\n"))
+                except (IOError, UnicodeDecodeError):
+                    pass
+
+        return total_files, total_lines