tarang 4.4.0__py3-none-any.whl

tarang/client/auth.py

@@ -0,0 +1,178 @@
+"""
+Tarang Authentication - CLI login and token management.
+
+Handles OAuth flow via browser and secure token storage.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import webbrowser
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from typing import Optional
+from urllib.parse import parse_qs, urlparse
+
+
+CONFIG_DIR = Path.home() / ".tarang"
+CONFIG_FILE = CONFIG_DIR / "config.json"
+
+
+class TarangAuth:
+    """
+    Handles CLI authentication via browser OAuth flow.
+
+    Stores credentials securely in ~/.tarang/config.json
+    """
+
+    def __init__(self, web_url: str = "https://devtarang.ai"):
+        self.web_url = web_url
+        self.token: Optional[str] = None
+
+    def load_credentials(self) -> Optional[dict]:
+        """Load saved credentials from config file."""
+        if CONFIG_FILE.exists():
+            try:
+                return json.loads(CONFIG_FILE.read_text())
+            except (json.JSONDecodeError, IOError):
+                return None
+        return None
+
+    def save_credentials(self, **kwargs) -> None:
+        """Save credentials to config file."""
+        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+        config = self.load_credentials() or {}
+        config.update(kwargs)
+        CONFIG_FILE.write_text(json.dumps(config, indent=2))
+        CONFIG_FILE.chmod(0o600)  # Secure permissions
+
+    def get_token(self) -> Optional[str]:
+        """Get saved auth token."""
+        creds = self.load_credentials()
+        return creds.get("token") if creds else None
+
+    def get_openrouter_key(self) -> Optional[str]:
+        """Get saved OpenRouter API key."""
+        creds = self.load_credentials()
+        return creds.get("openrouter_key") if creds else None
+
+    def save_token(self, token: str) -> None:
+        """Save auth token."""
+        self.save_credentials(token=token)
+        self.token = token
+
+    def save_openrouter_key(self, key: str) -> None:
+        """Save OpenRouter API key."""
+        self.save_credentials(openrouter_key=key)
+
+    def clear_credentials(self) -> None:
+        """Clear all saved credentials."""
+        if CONFIG_FILE.exists():
+            CONFIG_FILE.unlink()
+
+    async def login(self, callback_port: int = 54321) -> str:
+        """
+        Start OAuth flow via browser.
+
+        Steps:
+        1. Open browser to auth page
+        2. Start local server to receive callback
+        3. Store token securely
+
+        Returns:
+            The auth token
+        """
+        # Create callback server
+        server = _CallbackServer(("localhost", callback_port))
+
+        # Open browser to auth URL
+        auth_url = f"{self.web_url}/auth/cli?callback=http://localhost:{callback_port}"
+        print(f"Opening browser for authentication...")
+        print(f"If browser doesn't open, visit: {auth_url}")
+        webbrowser.open(auth_url)
+
+        print("\nWaiting for authentication...")
+        print("Please log in with GitHub in your browser.")
+
+        # Wait for callback (timeout 5 min)
+        try:
+            token = await asyncio.wait_for(
+                server.wait_for_token(),
+                timeout=300
+            )
+        except asyncio.TimeoutError:
+            raise TimeoutError("Authentication timed out. Please try again.")
+
+        # Save and return token
+        self.save_token(token)
+        return token
+
+    def is_authenticated(self) -> bool:
+        """Check if user is authenticated."""
+        return bool(self.get_token())
+
+    def has_openrouter_key(self) -> bool:
+        """Check if OpenRouter key is configured."""
+        return bool(self.get_openrouter_key())
+
+
+class _CallbackServer:
+    """Local HTTP server to receive OAuth callback."""
+
+    def __init__(self, address):
+        self.token: Optional[str] = None
+        self._received = asyncio.Event()
+        self.server = HTTPServer(address, self._make_handler())
+        self.server.timeout = 1  # Allow checking for cancellation
+
+    def _make_handler(self):
+        parent = self
+
+        class Handler(BaseHTTPRequestHandler):
+            def do_GET(self):
+                query = parse_qs(urlparse(self.path).query)
+                parent.token = query.get("token", [None])[0]
+
+                self.send_response(200)
+                self.send_header("Content-type", "text/html")
+                self.end_headers()
+
+                html = """
+                <!DOCTYPE html>
+                <html>
+                <head>
+                    <title>Tarang - Authentication Successful</title>
+                    <style>
+                        body { font-family: -apple-system, sans-serif; text-align: center; padding-top: 50px; }
+                        h1 { color: #10B981; }
+                    </style>
+                </head>
+                <body>
+                    <h1>Authentication Successful!</h1>
+                    <p>You can close this window and return to the terminal.</p>
+                </body>
+                </html>
+                """
+                self.wfile.write(html.encode())
+
+                if parent.token:
+                    print("\nReceived CLI callback, completing login...")
+                    parent._received.set()
+                else:
+                    print("\nReceived CLI callback without token. Please retry login.")
+
+            def log_message(self, *args):
+                pass  # Suppress HTTP logs
+
+        return Handler
+
+    async def wait_for_token(self) -> str:
+        """Wait for token from callback."""
+        loop = asyncio.get_event_loop()
+
+        while not self._received.is_set():
+            # Handle one request (non-blocking)
+            await loop.run_in_executor(None, self.server.handle_request)
+            await asyncio.sleep(0.1)
+
+        return self.token

tarang/context/__init__.py

@@ -0,0 +1,41 @@
+"""
+Tarang Context - Code indexing and retrieval system.
+
+Provides:
+- AST-based code chunking (tree-sitter)
+- BM25 keyword search
+- Symbol Graph (Code Knowledge Graph)
+- Graph-augmented retrieval
+"""
+
+from tarang.context.skeleton import SkeletonGenerator, ProjectSkeleton
+from tarang.context.chunker import Chunk, Chunker, SymbolInfo
+from tarang.context.bm25 import BM25Index, SearchResult
+from tarang.context.graph import SymbolGraph, SymbolNode
+from tarang.context.retriever import ContextRetriever, RetrievalResult, create_retriever
+from tarang.context.indexer import ProjectIndexer, IndexStats, index_project, get_retriever
+
+__all__ = [
+    # Skeleton (existing)
+    "SkeletonGenerator",
+    "ProjectSkeleton",
+    # Chunker
+    "Chunk",
+    "Chunker",
+    "SymbolInfo",
+    # BM25
+    "BM25Index",
+    "SearchResult",
+    # Graph
+    "SymbolGraph",
+    "SymbolNode",
+    # Retriever
+    "ContextRetriever",
+    "RetrievalResult",
+    "create_retriever",
+    # Indexer
+    "ProjectIndexer",
+    "IndexStats",
+    "index_project",
+    "get_retriever",
+]

tarang/context/bm25.py

@@ -0,0 +1,218 @@
+"""
+BM25 Index - Fast keyword-based search over code chunks.
+
+Uses the Okapi BM25 algorithm for ranking code chunks by relevance.
+"""
+from __future__ import annotations
+
+import pickle
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from .chunker import Chunk
+
+
+@dataclass
+class SearchResult:
+    """A search result with chunk and score."""
+    chunk: Chunk
+    score: float
+
+    def __repr__(self) -> str:
+        return f"SearchResult({self.chunk.id}, score={self.score:.3f})"
+
+
+class BM25Index:
+    """
+    BM25 index for code chunk search.
+
+    Uses rank-bm25 library for efficient BM25 ranking.
+    """
+
+    def __init__(self):
+        self._bm25 = None
+        self._chunks: List[Chunk] = []
+        self._chunk_map: Dict[str, Chunk] = {}  # id -> chunk
+
+    @property
+    def is_empty(self) -> bool:
+        """Check if index is empty."""
+        return len(self._chunks) == 0
+
+    def build(self, chunks: List[Chunk]) -> None:
+        """
+        Build BM25 index from chunks.
+
+        Args:
+            chunks: List of code chunks with tokens
+        """
+        from rank_bm25 import BM25Okapi
+
+        self._chunks = chunks
+        self._chunk_map = {c.id: c for c in chunks}
+
+        # Build BM25 index from tokens
+        tokenized_corpus = [c.tokens for c in chunks]
+        self._bm25 = BM25Okapi(tokenized_corpus)
+
+    def add_chunks(self, new_chunks: List[Chunk]) -> None:
+        """
+        Add chunks to existing index (requires rebuild).
+
+        For now, this rebuilds the entire index.
+        Future optimization: incremental BM25.
+        """
+        # Merge chunks, replacing existing by ID
+        for chunk in new_chunks:
+            self._chunk_map[chunk.id] = chunk
+
+        self._chunks = list(self._chunk_map.values())
+        self.build(self._chunks)
+
+    def remove_chunks(self, chunk_ids: List[str]) -> None:
+        """
+        Remove chunks from index (requires rebuild).
+        """
+        for chunk_id in chunk_ids:
+            self._chunk_map.pop(chunk_id, None)
+
+        self._chunks = list(self._chunk_map.values())
+        if self._chunks:
+            self.build(self._chunks)
+        else:
+            self._bm25 = None
+
+    def search(self, query: str, k: int = 10) -> List[SearchResult]:
+        """
+        Search for chunks matching query.
+
+        Args:
+            query: Search query (natural language)
+            k: Maximum results to return
+
+        Returns:
+            List of SearchResult sorted by score (descending)
+        """
+        if self._bm25 is None or not self._chunks:
+            return []
+
+        # Tokenize query
+        query_tokens = self._tokenize_query(query)
+        if not query_tokens:
+            return []
+
+        # Get BM25 scores
+        scores = self._bm25.get_scores(query_tokens)
+
+        # Pair with chunks and sort
+        results = [
+            SearchResult(chunk=self._chunks[i], score=score)
+            for i, score in enumerate(scores)
+            if score > 0
+        ]
+
+        results.sort(key=lambda r: r.score, reverse=True)
+        return results[:k]
+
+    def get_chunk(self, chunk_id: str) -> Optional[Chunk]:
+        """Get a chunk by ID."""
+        return self._chunk_map.get(chunk_id)
+
+    def get_chunks_for_file(self, file_path: str) -> List[Chunk]:
+        """Get all chunks for a file."""
+        return [c for c in self._chunks if c.file == file_path]
+
+    def _tokenize_query(self, query: str) -> List[str]:
+        """Tokenize search query."""
+        # Split on whitespace and punctuation
+        words = re.findall(r'\b\w+\b', query.lower())
+
+        tokens = []
+        for word in words:
+            # Split snake_case
+            if "_" in word:
+                tokens.extend(word.split("_"))
+            # Split camelCase
+            elif any(c.isupper() for c in word[1:]):
+                parts = re.findall(r'[A-Z]?[a-z]+|[A-Z]+(?=[A-Z]|$)', word)
+                tokens.extend(p.lower() for p in parts)
+            else:
+                tokens.append(word)
+
+        # Filter stop words
+        stop_words = {
+            "the", "a", "an", "is", "are", "was", "were", "be", "been",
+            "have", "has", "had", "do", "does", "did", "will", "would",
+            "could", "should", "may", "might", "must", "can", "need",
+            "to", "of", "in", "for", "on", "with", "at", "by", "from",
+            "as", "into", "through", "during", "before", "after",
+            "and", "but", "if", "or", "because", "until", "while",
+            "this", "that", "these", "those", "i", "me", "my", "we",
+            "you", "your", "it", "its", "they", "them", "their",
+            "what", "which", "who", "how", "where", "when", "why",
+            "add", "create", "make", "build", "implement", "write",
+            "fix", "update", "change", "modify", "remove", "delete",
+        }
+
+        return [t for t in tokens if len(t) > 2 and t not in stop_words]
+
+    def save(self, path: Path) -> None:
+        """
+        Save index to disk.
+
+        Args:
+            path: Path to save file (pickle format)
+        """
+        data = {
+            "chunks": [c.to_dict() for c in self._chunks],
+            "bm25": self._bm25,
+        }
+        with open(path, "wb") as f:
+            pickle.dump(data, f)
+
+    def load(self, path: Path) -> bool:
+        """
+        Load index from disk.
+
+        Args:
+            path: Path to saved file
+
+        Returns:
+            True if loaded successfully
+        """
+        if not path.exists():
+            return False
+
+        try:
+            with open(path, "rb") as f:
+                data = pickle.load(f)
+
+            self._chunks = [Chunk.from_dict(d) for d in data["chunks"]]
+            self._chunk_map = {c.id: c for c in self._chunks}
+            self._bm25 = data["bm25"]
+            return True
+
+        except Exception:
+            return False
+
+    def stats(self) -> Dict:
+        """Get index statistics."""
+        if not self._chunks:
+            return {
+                "total_chunks": 0,
+                "total_files": 0,
+                "chunk_types": {},
+            }
+
+        files = set(c.file for c in self._chunks)
+        types = {}
+        for c in self._chunks:
+            types[c.type] = types.get(c.type, 0) + 1
+
+        return {
+            "total_chunks": len(self._chunks),
+            "total_files": len(files),
+            "chunk_types": types,
+        }