mita-code 0.1.0__py3-none-any.whl

mita/index/manager.py

@@ -0,0 +1,170 @@
+"""CLI command handlers for codebase indexing."""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+from rich.console import Console
+from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn
+from rich.syntax import Syntax
+from rich.table import Table
+
+from mita.config.loader import load_config
+from mita.index.embeddings import EmbeddingClient
+from mita.index.parser import parse_codebase
+from mita.index.retriever import Retriever
+from mita.index.store import IndexStore
+
+console = Console()
+
+
+def _get_index_dir() -> Path:
+    return Path.cwd() / ".mita" / "index"
+
+
+async def build_index(force: bool = False, pull_model_fn: object | None = None) -> None:
+    """Build the codebase index.
+
+    Args:
+        force: Rebuild even if index already exists.
+        pull_model_fn: Async callback to pull the embedding model if missing.
+            Signature: async (MitaConfig) -> bool. Called from the CLI layer
+            so it can import from mita.models without violating dependency rules.
+    """
+    config = load_config()
+    index_dir = _get_index_dir()
+    store = IndexStore(index_dir)
+
+    # Check if index already exists
+    if store.exists() and not force:
+        console.print("[yellow]Index already exists. Use --force to rebuild.[/yellow]")
+        return
+
+    # Check embedding model availability
+    embedder = EmbeddingClient(config)
+    if not await embedder.is_model_available():
+        if pull_model_fn is not None:
+            if not await pull_model_fn(config):  # type: ignore[operator]
+                return
+        else:
+            console.print("[red]Embedding model not available.[/red]")
+            return
+
+    start_time = time.monotonic()
+
+    # Parse codebase
+    root = Path.cwd()
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        console=console,
+    ) as progress:
+        progress.add_task("Parsing codebase...", total=None)
+        chunks = parse_codebase(root, config.index)
+
+    if not chunks:
+        console.print("[yellow]No code chunks found to index.[/yellow]")
+        return
+
+    console.print(f"Found {len(chunks)} chunks from codebase.")
+
+    # Generate embeddings
+    texts = [c.content for c in chunks]
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TextColumn("{task.completed}/{task.total}"),
+        console=console,
+    ) as progress:
+        task = progress.add_task("Generating embeddings...", total=len(texts))
+        embeddings = []
+        batch_size = 32
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
+            batch_embeddings = await embedder.embed_texts(batch)
+            embeddings.extend(batch_embeddings)
+            progress.update(task, completed=min(i + batch_size, len(texts)))
+
+    # Attach embeddings to chunks
+    for chunk, embedding in zip(chunks, embeddings):
+        chunk.embedding = embedding
+
+    # Store in LanceDB
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        console=console,
+    ) as progress:
+        progress.add_task("Writing index...", total=None)
+        await store.create_or_replace(chunks)
+
+    elapsed = time.monotonic() - start_time
+    file_count = len({c.file_path for c in chunks})
+    console.print(
+        f"[green]Indexed {len(chunks)} chunks from {file_count} files in {elapsed:.1f}s.[/green]"
+    )
+
+
+async def show_index_status() -> None:
+    """Show index statistics."""
+    config = load_config()
+    store = IndexStore(_get_index_dir())
+    stats = await store.status()
+
+    if not stats["exists"]:
+        console.print("[yellow]No index found. Run 'mita index build' first.[/yellow]")
+        return
+
+    table = Table(title="Index Status")
+    table.add_column("Metric", style="bold")
+    table.add_column("Value")
+    table.add_row("Chunks", str(stats["chunks"]))
+    table.add_row("Files", str(stats["files"]))
+    table.add_row("Location", str(_get_index_dir()))
+    table.add_row("Embedding model", config.model.embedding)
+    console.print(table)
+
+
+async def search_index(query: str, top_k: int = 10) -> None:
+    """Search the index and display results."""
+    config = load_config()
+    retriever = Retriever(config, index_dir=_get_index_dir())
+
+    if not retriever.is_available():
+        console.print("[yellow]No index found. Run 'mita index build' first.[/yellow]")
+        return
+
+    results = await retriever.retrieve(query, top_k=top_k)
+
+    if not results:
+        console.print("[yellow]No results found.[/yellow]")
+        return
+
+    console.print(f"[bold]Found {len(results)} results:[/bold]\n")
+    for i, r in enumerate(results, 1):
+        c = r.chunk
+        header = f"{i}. {c.file_path}:{c.start_line}-{c.end_line}"
+        if c.symbol:
+            header += f" ({c.symbol})"
+        header += f"  [dim]score: {r.score:.3f}[/dim]"
+        console.print(header)
+        console.print(Syntax(c.content, c.language, line_numbers=True, start_line=c.start_line))
+        console.print()
+
+
+async def clear_index() -> None:
+    """Clear the index."""
+    store = IndexStore(_get_index_dir())
+
+    if not store.exists():
+        console.print("[yellow]No index found.[/yellow]")
+        return
+
+    confirm = console.input("Delete the index? [y/N] ").strip().lower()
+    if confirm in ("y", "yes"):
+        await store.clear()
+        console.print("[green]Index cleared.[/green]")
+    else:
+        console.print("Cancelled.")

mita/index/parser.py

@@ -0,0 +1,331 @@
+"""Tree-sitter code parsing and chunking for codebase indexing."""
+
+from __future__ import annotations
+
+from fnmatch import fnmatch
+from pathlib import Path
+from typing import Any
+
+from mita.config.schema import IndexSettings
+from mita.index.store import CodeChunk
+
+# Extension → tree-sitter language name
+EXTENSION_MAP: dict[str, str] = {
+    ".py": "python",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".ts": "typescript",
+    ".tsx": "tsx",
+    ".go": "go",
+    ".rs": "rust",
+    ".java": "java",
+    ".c": "c",
+    ".h": "c",
+    ".cpp": "cpp",
+    ".hpp": "cpp",
+    ".cc": "cpp",
+    ".rb": "ruby",
+    ".php": "php",
+    ".swift": "swift",
+    ".kt": "kotlin",
+    ".scala": "scala",
+    ".lua": "lua",
+    ".sh": "bash",
+    ".bash": "bash",
+    ".zsh": "bash",
+    ".css": "css",
+    ".html": "html",
+    ".json": "json",
+    ".yaml": "yaml",
+    ".yml": "yaml",
+    ".toml": "toml",
+}
+
+# Node types that represent semantic code units per language
+CHUNK_NODE_TYPES: dict[str, set[str]] = {
+    "python": {"function_definition", "class_definition", "decorated_definition"},
+    "javascript": {
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "export_statement",
+    },
+    "typescript": {
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "export_statement",
+        "interface_declaration",
+        "type_alias_declaration",
+    },
+    "tsx": {
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "export_statement",
+        "interface_declaration",
+        "type_alias_declaration",
+    },
+    "go": {"function_declaration", "method_declaration", "type_declaration"},
+    "rust": {"function_item", "impl_item", "struct_item", "enum_item", "trait_item"},
+    "java": {"class_declaration", "method_declaration", "interface_declaration"},
+    "c": {"function_definition", "struct_specifier"},
+    "cpp": {"function_definition", "class_specifier", "struct_specifier"},
+    "ruby": {"method", "class", "module"},
+}
+
+
+def parse_codebase(root: Path, config: IndexSettings) -> list[CodeChunk]:
+    """Parse all supported files under root into code chunks."""
+    chunks: list[CodeChunk] = []
+    for file_path in _discover_files(root, config.exclude_patterns):
+        file_chunks = parse_file(file_path, root, config)
+        chunks.extend(file_chunks)
+    return chunks
+
+
+def parse_file(file_path: Path, root: Path, config: IndexSettings) -> list[CodeChunk]:
+    """Parse a single file into code chunks."""
+    if _is_binary(file_path):
+        return []
+
+    try:
+        content = file_path.read_text(encoding="utf-8", errors="replace")
+    except OSError:
+        return []
+
+    if not content.strip():
+        return []
+
+    rel_path = str(file_path.relative_to(root))
+    language = get_language_for_file(file_path)
+
+    if language and language in CHUNK_NODE_TYPES:
+        chunks = _parse_with_treesitter(content, rel_path, language, config)
+        if chunks:
+            return chunks
+
+    # Fallback: line-based chunking
+    return _chunk_by_lines(content, rel_path, language or "text", config)
+
+
+def get_language_for_file(path: Path) -> str | None:
+    """Map a file extension to a tree-sitter language name."""
+    return EXTENSION_MAP.get(path.suffix.lower())
+
+
+def _discover_files(root: Path, exclude_patterns: list[str]) -> list[Path]:
+    """Walk the project and yield non-excluded files."""
+    files = []
+    for path in sorted(root.rglob("*")):
+        if not path.is_file():
+            continue
+        rel = str(path.relative_to(root))
+        if _matches_exclude(rel, exclude_patterns):
+            continue
+        files.append(path)
+    return files
+
+
+def _matches_exclude(rel_path: str, patterns: list[str]) -> bool:
+    """Check if a relative path matches any exclude pattern."""
+    for pattern in patterns:
+        if fnmatch(rel_path, pattern):
+            return True
+        # Also check each path component for directory patterns
+        parts = rel_path.split("/")
+        for i in range(len(parts)):
+            partial = "/".join(parts[: i + 1])
+            if fnmatch(partial, pattern.rstrip("/")):
+                return True
+            if fnmatch(partial + "/", pattern):
+                return True
+    return False
+
+
+def _is_binary(path: Path) -> bool:
+    """Heuristic: file is binary if first 8KB contains null bytes."""
+    try:
+        with open(path, "rb") as f:
+            chunk = f.read(8192)
+        return b"\x00" in chunk
+    except OSError:
+        return True
+
+
+def _parse_with_treesitter(
+    content: str,
+    rel_path: str,
+    language: str,
+    config: IndexSettings,
+) -> list[CodeChunk]:
+    """Parse a file using tree-sitter and extract semantic chunks."""
+    try:
+        from tree_sitter_language_pack import get_parser
+    except ImportError:
+        return []
+
+    try:
+        parser = get_parser(language)  # type: ignore[arg-type]
+    except Exception:  # noqa: BLE001
+        return []
+
+    tree = parser.parse(content.encode("utf-8"))
+    node_types = CHUNK_NODE_TYPES.get(language, set())
+    lines = content.split("\n")
+    chunks: list[CodeChunk] = []
+
+    for node in _walk_top_level(tree.root_node):
+        if node.type not in node_types:
+            continue
+
+        start_line = node.start_point[0] + 1  # 1-indexed
+        end_line = node.end_point[0] + 1
+        node_content = "\n".join(lines[start_line - 1 : end_line])
+        symbol = _extract_symbol(node, language)
+
+        # Split large nodes
+        char_budget = config.chunk_size * 4  # ~4 chars per token
+        if len(node_content) > char_budget:
+            sub_chunks = _split_large_content(
+                node_content, start_line, rel_path, language, symbol, config
+            )
+            chunks.extend(sub_chunks)
+        else:
+            chunks.append(
+                CodeChunk(
+                    file_path=rel_path,
+                    start_line=start_line,
+                    end_line=end_line,
+                    content=node_content,
+                    language=language,
+                    symbol=symbol,
+                )
+            )
+
+    return chunks
+
+
+def _walk_top_level(node: Any) -> list[Any]:
+    """Get top-level children of the root node (non-recursive)."""
+    children = getattr(node, "children", [])
+    return list(children)
+
+
+def _extract_symbol(node: Any, language: str) -> str | None:
+    """Extract the name of a function/class node."""
+    # Most languages store the name in a child node of type "identifier" or "name"
+    for child in getattr(node, "children", []):
+        child_type = getattr(child, "type", "")
+        if child_type in ("identifier", "name", "property_identifier"):
+            return getattr(child, "text", b"").decode("utf-8", errors="replace")
+    return None
+
+
+def _split_large_content(
+    content: str,
+    base_start_line: int,
+    rel_path: str,
+    language: str,
+    symbol: str | None,
+    config: IndexSettings,
+) -> list[CodeChunk]:
+    """Split a large chunk into overlapping sub-chunks."""
+    lines = content.split("\n")
+    char_budget = config.chunk_size * 4
+    overlap_chars = config.chunk_overlap * 4
+    chunks: list[CodeChunk] = []
+
+    start = 0
+    while start < len(lines):
+        # Accumulate lines until we hit the budget
+        end = start
+        total_chars = 0
+        while end < len(lines) and total_chars < char_budget:
+            total_chars += len(lines[end]) + 1
+            end += 1
+
+        chunk_content = "\n".join(lines[start:end])
+        chunks.append(
+            CodeChunk(
+                file_path=rel_path,
+                start_line=base_start_line + start,
+                end_line=base_start_line + end - 1,
+                content=chunk_content,
+                language=language,
+                symbol=symbol,
+            )
+        )
+
+        # If we consumed all lines, we're done
+        if end >= len(lines):
+            break
+
+        # Advance with overlap
+        overlap_lines = 0
+        overlap_total = 0
+        for i in range(end - 1, start, -1):
+            overlap_total += len(lines[i]) + 1
+            overlap_lines += 1
+            if overlap_total >= overlap_chars:
+                break
+
+        new_start = end - overlap_lines
+        if new_start <= start:
+            break
+        start = new_start
+
+    return chunks
+
+
+def _chunk_by_lines(
+    content: str,
+    rel_path: str,
+    language: str,
+    config: IndexSettings,
+) -> list[CodeChunk]:
+    """Fallback: chunk file by line windows."""
+    lines = content.split("\n")
+    char_budget = config.chunk_size * 4
+    overlap_chars = config.chunk_overlap * 4
+    chunks: list[CodeChunk] = []
+
+    start = 0
+    while start < len(lines):
+        end = start
+        total_chars = 0
+        while end < len(lines) and total_chars < char_budget:
+            total_chars += len(lines[end]) + 1
+            end += 1
+
+        chunk_content = "\n".join(lines[start:end])
+        if chunk_content.strip():
+            chunks.append(
+                CodeChunk(
+                    file_path=rel_path,
+                    start_line=start + 1,
+                    end_line=end,
+                    content=chunk_content,
+                    language=language,
+                )
+            )
+
+        # If we consumed all lines, we're done
+        if end >= len(lines):
+            break
+
+        # Advance with overlap
+        overlap_lines = 0
+        overlap_total = 0
+        for i in range(end - 1, start, -1):
+            overlap_total += len(lines[i]) + 1
+            overlap_lines += 1
+            if overlap_total >= overlap_chars:
+                break
+
+        new_start = end - overlap_lines
+        if new_start <= start:
+            break
+        start = new_start
+
+    return chunks

mita/index/retriever.py

@@ -0,0 +1,53 @@
+"""RAG retrieval pipeline: vector search over indexed code chunks."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from mita.config.schema import MitaConfig
+from mita.index.embeddings import EmbeddingClient
+from mita.index.store import IndexStore, SearchResult
+
+
+def _default_index_dir() -> Path:
+    """Determine the index directory for the current project."""
+    return Path.cwd() / ".mita" / "index"
+
+
+class Retriever:
+    """High-level retrieval combining embeddings + vector search."""
+
+    def __init__(self, config: MitaConfig, index_dir: Path | None = None) -> None:
+        self._config = config
+        self._index_dir = index_dir or _default_index_dir()
+        self._store = IndexStore(self._index_dir)
+        self._embedder = EmbeddingClient(config)
+
+    async def retrieve(self, query: str, top_k: int | None = None) -> list[SearchResult]:
+        """Retrieve relevant code chunks for a query."""
+        k = top_k or self._config.index.top_k
+        query_embedding = await self._embedder.embed_single(query)
+        return await self._store.search(query_embedding, top_k=k)
+
+    async def retrieve_formatted(self, query: str, top_k: int | None = None) -> str:
+        """Retrieve and format results as a context string for the LLM."""
+        results = await self.retrieve(query, top_k=top_k)
+        if not results:
+            return ""
+        return format_results(results)
+
+    def is_available(self) -> bool:
+        """Check if the index exists and is ready to query."""
+        return self._store.exists()
+
+
+def format_results(results: list[SearchResult]) -> str:
+    """Format search results as a readable context block."""
+    parts: list[str] = []
+    for r in results:
+        c = r.chunk
+        header = f"## {c.file_path}:{c.start_line}-{c.end_line}"
+        if c.symbol:
+            header += f" ({c.symbol})"
+        parts.append(f"{header}\n```{c.language}\n{c.content}\n```")
+    return "\n\n".join(parts)

mita/index/store.py

@@ -0,0 +1,143 @@
+"""LanceDB vector store operations for code chunk storage and retrieval."""
+
+from __future__ import annotations
+
+import asyncio
+import shutil
+from pathlib import Path
+from typing import Any
+
+import lancedb  # type: ignore[import-untyped]
+from pydantic import BaseModel, Field
+
+
+class CodeChunk(BaseModel):
+    """A chunk of code extracted from a source file."""
+
+    file_path: str
+    start_line: int
+    end_line: int
+    content: str
+    language: str
+    symbol: str | None = None
+    embedding: list[float] | None = Field(default=None, exclude=True)
+
+
+class SearchResult(BaseModel):
+    """A code chunk with its similarity score."""
+
+    chunk: CodeChunk
+    score: float
+
+
+class IndexStore:
+    """LanceDB-backed vector store for code chunks."""
+
+    TABLE_NAME = "code_chunks"
+
+    def __init__(self, index_dir: Path) -> None:
+        self._index_dir = index_dir
+        self._db_path = index_dir / "lancedb"
+
+    def _connect(self) -> lancedb.DBConnection:
+        self._db_path.parent.mkdir(parents=True, exist_ok=True)
+        return lancedb.connect(str(self._db_path))
+
+    async def create_or_replace(self, chunks: list[CodeChunk]) -> None:
+        """Create or replace the index with the given chunks."""
+        if not chunks:
+            return
+
+        def _sync() -> None:
+            db = self._connect()
+            records = _chunks_to_records(chunks)
+            if self.TABLE_NAME in db.table_names():
+                db.drop_table(self.TABLE_NAME)
+            db.create_table(self.TABLE_NAME, data=records)
+
+        await asyncio.to_thread(_sync)
+
+    async def search(self, query_embedding: list[float], top_k: int = 10) -> list[SearchResult]:
+        """Vector similarity search."""
+
+        def _sync() -> list[SearchResult]:
+            db = self._connect()
+            if self.TABLE_NAME not in db.table_names():
+                return []
+            table = db.open_table(self.TABLE_NAME)
+            results = table.search(query_embedding).limit(top_k).to_pandas()
+
+            search_results: list[SearchResult] = []
+            for _, row in results.iterrows():
+                distance = row.get("_distance", 0.0)
+                score = 1.0 / (1.0 + distance)
+                chunk = CodeChunk(
+                    file_path=row["file_path"],
+                    start_line=int(row["start_line"]),
+                    end_line=int(row["end_line"]),
+                    content=row["content"],
+                    language=row["language"],
+                    symbol=row.get("symbol"),
+                )
+                search_results.append(SearchResult(chunk=chunk, score=score))
+            return search_results
+
+        return await asyncio.to_thread(_sync)
+
+    async def clear(self) -> None:
+        """Delete the entire index."""
+
+        def _sync() -> None:
+            if self._db_path.exists():
+                shutil.rmtree(self._db_path)
+
+        await asyncio.to_thread(_sync)
+
+    async def status(self) -> dict[str, Any]:
+        """Return index statistics."""
+
+        def _sync() -> dict[str, Any]:
+            if not self._db_path.exists():
+                return {"exists": False, "chunks": 0, "files": 0}
+            db = self._connect()
+            if self.TABLE_NAME not in db.table_names():
+                return {"exists": False, "chunks": 0, "files": 0}
+            table = db.open_table(self.TABLE_NAME)
+            df = table.to_pandas()
+            return {
+                "exists": True,
+                "chunks": len(df),
+                "files": df["file_path"].nunique() if len(df) > 0 else 0,
+            }
+
+        return await asyncio.to_thread(_sync)
+
+    def exists(self) -> bool:
+        """Check if the index exists on disk."""
+        if not self._db_path.exists():
+            return False
+        try:
+            db = self._connect()
+            return self.TABLE_NAME in db.table_names()
+        except Exception:  # noqa: BLE001
+            return False
+
+
+def _chunks_to_records(chunks: list[CodeChunk]) -> list[dict[str, Any]]:
+    """Convert CodeChunks (with embeddings) to LanceDB-compatible records."""
+    records = []
+    for chunk in chunks:
+        if chunk.embedding is None:
+            continue
+        records.append(
+            {
+                "file_path": chunk.file_path,
+                "start_line": chunk.start_line,
+                "end_line": chunk.end_line,
+                "content": chunk.content,
+                "language": chunk.language,
+                "symbol": chunk.symbol or "",
+                "vector": chunk.embedding,
+            }
+        )
+    return records

mita/llm/__init__.py

@@ -0,0 +1 @@
+"""LLM client abstraction: LiteLLM + Instructor + streaming."""