mita-code 0.1.0__py3-none-any.whl

mita/models/recommender.py

@@ -0,0 +1,88 @@
+"""Filter and rank models by hardware capability."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from mita.models.hardware import HardwareInfo
+from mita.models.registry import ModelCard, get_coding_models
+
+
+class ModelRecommendation(BaseModel):
+    """A model recommendation with a fit score and notes."""
+
+    model: ModelCard
+    fit_score: float  # 0.0–1.0
+    notes: str
+
+
+def recommend_models(hw: HardwareInfo) -> list[ModelRecommendation]:
+    """Recommend coding models that fit the given hardware.
+
+    Returns a list sorted by fit_score descending (best fit first).
+    Only includes models that can actually run on the hardware.
+    """
+    available_vram = hw.available_vram_gb
+    recommendations: list[ModelRecommendation] = []
+
+    for model in get_coding_models():
+        if model.min_vram_gb <= 0:
+            continue
+        if model.min_vram_gb > available_vram:
+            continue
+        if model.min_ram_gb > hw.ram_gb:
+            continue
+
+        fit_score, notes = _score_model(model, hw, available_vram)
+        recommendations.append(ModelRecommendation(model=model, fit_score=fit_score, notes=notes))
+
+    recommendations.sort(key=lambda r: r.fit_score, reverse=True)
+    return recommendations
+
+
+def _score_model(model: ModelCard, hw: HardwareInfo, available_vram: float) -> tuple[float, str]:
+    """Score a model's fit for the hardware. Returns (score, notes)."""
+    notes_parts: list[str] = []
+
+    # Base score: how much headroom do we have?
+    # Ratio of available VRAM to required VRAM
+    headroom_ratio = available_vram / model.min_vram_gb
+    if headroom_ratio >= 2.0:
+        vram_score = 1.0
+        notes_parts.append(f"Plenty of headroom ({available_vram:.0f}GB available)")
+    elif headroom_ratio >= 1.5:
+        vram_score = 0.9
+        notes_parts.append(f"Comfortable fit ({available_vram:.0f}GB available)")
+    elif headroom_ratio >= 1.2:
+        vram_score = 0.7
+        notes_parts.append(f"Fits with some room ({available_vram:.0f}GB available)")
+    else:
+        vram_score = 0.5
+        notes_parts.append(f"Tight fit — uses ~{model.min_vram_gb:.0f}GB of {available_vram:.0f}GB")
+
+    # Bonus for tool call support (important for agentic use)
+    tool_bonus = 0.1 if model.tool_call_support else 0.0
+
+    # Bonus for being a recommended model
+    recommended_bonus = 0.05 if "recommended" in model.tags else 0.0
+
+    # Penalty for Apple Silicon running very large models (thermal throttling)
+    thermal_penalty = 0.0
+    if hw.apple_silicon and model.min_vram_gb > hw.ram_gb * 0.5:
+        thermal_penalty = 0.1
+        notes_parts.append("May cause thermal throttling on Apple Silicon")
+
+    # Context window bonus (larger is better for agentic coding)
+    ctx_bonus = min(0.05, model.context_window / 1_000_000)
+
+    score = min(1.0, vram_score + tool_bonus + recommended_bonus + ctx_bonus - thermal_penalty)
+
+    if model.tool_call_support:
+        notes_parts.append("Supports tool calls")
+    else:
+        notes_parts.append("No tool call support — will use JSON fallback")
+
+    if hw.unified_memory:
+        notes_parts.append("Unified memory — shared with system")
+
+    return round(score, 2), ". ".join(notes_parts)

mita/models/registry.py

@@ -0,0 +1,167 @@
+"""Curated registry of coding-focused models for Ollama."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class ModelCard(BaseModel):
+    """Metadata for a coding-focused model in the curated registry."""
+
+    name: str  # Ollama model tag, e.g. "qwen2.5-coder:7b"
+    family: str  # "qwen", "codellama", "deepseek", etc.
+    param_count: str  # "7B", "14B", "32B"
+    min_ram_gb: float
+    min_vram_gb: float
+    context_window: int
+    quantization: str = "Q4_K_M"
+    tool_call_support: bool = True
+    description: str = ""
+    tags: list[str] = Field(default_factory=list)
+
+
+# Curated list of coding models that work well with Ollama
+CODING_MODELS: list[ModelCard] = [
+    # ── Small (< 4GB) ────────────────────────────────────────────
+    ModelCard(
+        name="qwen2.5-coder:1.5b",
+        family="qwen",
+        param_count="1.5B",
+        min_ram_gb=2,
+        min_vram_gb=2,
+        context_window=32768,
+        description="Tiny but capable code model. Good for quick completions.",
+        tags=["code", "fast", "small"],
+        tool_call_support=False,
+    ),
+    ModelCard(
+        name="qwen2.5-coder:3b",
+        family="qwen",
+        param_count="3B",
+        min_ram_gb=3,
+        min_vram_gb=3,
+        context_window=32768,
+        description="Small code model. Good balance of speed and quality.",
+        tags=["code", "fast"],
+        tool_call_support=False,
+    ),
+    ModelCard(
+        name="deepseek-coder-v2:lite",
+        family="deepseek",
+        param_count="2.4B",
+        min_ram_gb=3,
+        min_vram_gb=2,
+        context_window=16384,
+        description="Lightweight DeepSeek coder variant.",
+        tags=["code", "fast", "small"],
+        tool_call_support=False,
+    ),
+    # ── Medium (4–8GB) ───────────────────────────────────────────
+    ModelCard(
+        name="qwen2.5-coder:7b",
+        family="qwen",
+        param_count="7B",
+        min_ram_gb=6,
+        min_vram_gb=5,
+        context_window=32768,
+        description="Strong code model. Best default for 8GB+ machines.",
+        tags=["code", "instruct", "recommended"],
+    ),
+    ModelCard(
+        name="codellama:7b",
+        family="codellama",
+        param_count="7B",
+        min_ram_gb=6,
+        min_vram_gb=5,
+        context_window=16384,
+        description="Meta's Code Llama. Solid general-purpose code model.",
+        tags=["code", "instruct"],
+    ),
+    ModelCard(
+        name="deepseek-coder-v2:16b",
+        family="deepseek",
+        param_count="16B",
+        min_ram_gb=10,
+        min_vram_gb=10,
+        context_window=65536,
+        description="DeepSeek Coder V2. Strong MoE architecture.",
+        tags=["code", "instruct"],
+    ),
+    # ── Large (8–16GB) ───────────────────────────────────────────
+    ModelCard(
+        name="qwen2.5-coder:14b",
+        family="qwen",
+        param_count="14B",
+        min_ram_gb=10,
+        min_vram_gb=10,
+        context_window=32768,
+        description="Excellent code model. Strong tool call support.",
+        tags=["code", "instruct", "recommended"],
+    ),
+    ModelCard(
+        name="codellama:13b",
+        family="codellama",
+        param_count="13B",
+        min_ram_gb=10,
+        min_vram_gb=10,
+        context_window=16384,
+        description="Larger Code Llama with better reasoning.",
+        tags=["code", "instruct"],
+    ),
+    # ── XL (16–32GB) ─────────────────────────────────────────────
+    ModelCard(
+        name="qwen2.5-coder:32b",
+        family="qwen",
+        param_count="32B",
+        min_ram_gb=20,
+        min_vram_gb=20,
+        context_window=32768,
+        description="Top-tier local code model. Excellent tool call reliability.",
+        tags=["code", "instruct", "recommended"],
+    ),
+    ModelCard(
+        name="codellama:34b",
+        family="codellama",
+        param_count="34B",
+        min_ram_gb=22,
+        min_vram_gb=22,
+        context_window=16384,
+        description="Largest Code Llama. Near-cloud quality.",
+        tags=["code", "instruct"],
+    ),
+    # ── Embedding models ─────────────────────────────────────────
+    ModelCard(
+        name="nomic-embed-text",
+        family="nomic",
+        param_count="137M",
+        min_ram_gb=1,
+        min_vram_gb=1,
+        context_window=8192,
+        description="Default embedding model for codebase indexing.",
+        tags=["embedding"],
+        tool_call_support=False,
+    ),
+]
+
+
+def get_registry() -> list[ModelCard]:
+    """Return the full curated model registry."""
+    return list(CODING_MODELS)
+
+
+def find_model(name: str) -> ModelCard | None:
+    """Find a model in the registry by name (exact match)."""
+    for model in CODING_MODELS:
+        if model.name == name:
+            return model
+    return None
+
+
+def get_coding_models() -> list[ModelCard]:
+    """Return only coding models (exclude embedding models)."""
+    return [m for m in CODING_MODELS if "embedding" not in m.tags]
+
+
+def get_embedding_models() -> list[ModelCard]:
+    """Return only embedding models."""
+    return [m for m in CODING_MODELS if "embedding" in m.tags]

mita/models/server.py

@@ -0,0 +1,262 @@
+"""Ollama server lifecycle management — start, stop, health check.
+
+Runs Ollama as a detached daemon tracked via a PID file so it persists
+across short-lived CLI commands like ``mita models list``.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import signal
+import subprocess
+import time
+from pathlib import Path
+
+from rich.console import Console
+
+from mita.models.ollama_client import OllamaClient
+
+
+def _get_pid_dir() -> Path:
+    """Return the directory for the PID file (lazy to avoid module-level Path.home())."""
+    return Path.home() / ".config" / "mita"
+
+
+def _get_pid_file() -> Path:
+    """Return the PID file path (lazy to avoid module-level Path.home())."""
+    return _get_pid_dir() / "ollama.pid"
+
+
+def find_ollama_binary() -> str | None:
+    """Locate the ollama binary on the system."""
+    return shutil.which("ollama")
+
+
+def is_server_running(host: str = "http://localhost:11434", timeout: int = 5) -> bool:
+    """Check if an Ollama server is reachable."""
+    client = OllamaClient(host=host, timeout=timeout)
+    return client.is_running()
+
+
+def _read_pid() -> int | None:
+    """Read the PID from the PID file, return None if missing or stale."""
+    try:
+        pid = int(_get_pid_file().read_text().strip())
+        # Check if process is alive
+        os.kill(pid, 0)
+        return pid
+    except (FileNotFoundError, ValueError, ProcessLookupError, PermissionError, OSError):
+        _remove_pid_file()
+        return None
+
+
+def _write_pid(pid: int) -> None:
+    """Write a PID to the PID file."""
+    _get_pid_dir().mkdir(parents=True, exist_ok=True)
+    _get_pid_file().write_text(str(pid))
+
+
+def _remove_pid_file() -> None:
+    """Remove the PID file if it exists."""
+    try:
+        _get_pid_file().unlink()
+    except FileNotFoundError:
+        pass
+
+
+def start_server(
+    host: str = "http://localhost:11434",
+    timeout: int = 30,
+    console: Console | None = None,
+) -> bool:
+    """Start the Ollama server as a detached daemon.
+
+    The server persists after Mita exits. Its PID is tracked in
+    ``~/.config/mita/ollama.pid`` so ``stop_server`` can find it later.
+
+    Returns True if the server is running (either already was or we started it).
+    """
+    if is_server_running(host, timeout=5):
+        return True
+
+    binary = find_ollama_binary()
+    if binary is None:
+        if console:
+            console.print(
+                "[red]Ollama binary not found.[/red]\n"
+                "Install Ollama from [bold]https://ollama.com[/bold]"
+            )
+        return False
+
+    if console:
+        console.print("[dim]Starting Ollama server...[/dim]")
+
+    env = os.environ.copy()
+    env_host = host.replace("http://", "").replace("https://", "")
+    env["OLLAMA_HOST"] = env_host
+
+    try:
+        proc = subprocess.Popen(
+            [binary, "serve"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            env=env,
+            start_new_session=True,
+        )
+    except OSError as e:
+        if console:
+            console.print(f"[red]Failed to start Ollama: {e}[/red]")
+        return False
+
+    _write_pid(proc.pid)
+
+    # Wait for the server to become healthy
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if proc.poll() is not None:
+            if console:
+                console.print("[red]Ollama server exited unexpectedly.[/red]")
+            _remove_pid_file()
+            return False
+        if is_server_running(host, timeout=2):
+            if console:
+                console.print("[green]Ollama server started.[/green]")
+            return True
+        time.sleep(0.5)
+
+    if console:
+        console.print("[red]Ollama server did not become ready in time.[/red]")
+    stop_server()
+    return False
+
+
+def stop_server(console: Console | None = None) -> bool:
+    """Stop the Ollama server if it was started by Mita (tracked via PID file).
+
+    Returns True if we stopped it, False if there was nothing to stop.
+    """
+    pid = _read_pid()
+    if pid is None:
+        return False
+
+    try:
+        os.kill(pid, signal.SIGTERM)
+        # Wait for graceful shutdown
+        deadline = time.monotonic() + 10
+        while time.monotonic() < deadline:
+            try:
+                os.kill(pid, 0)
+                time.sleep(0.25)
+            except ProcessLookupError:
+                break
+        else:
+            # Still alive after 10s — force kill
+            try:
+                os.kill(pid, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+    except ProcessLookupError:
+        pass  # Already dead
+
+    _remove_pid_file()
+
+    if console:
+        console.print("[dim]Ollama server stopped.[/dim]")
+
+    return True
+
+
+def is_managed() -> bool:
+    """Return True if a Mita-started Ollama server is tracked and alive."""
+    return _read_pid() is not None
+
+
+def ensure_server(
+    host: str = "http://localhost:11434",
+    auto_manage: bool = True,
+    console: Console | None = None,
+) -> bool:
+    """Ensure Ollama is available — start it if auto_manage is enabled."""
+    if is_server_running(host, timeout=5):
+        return True
+
+    if not auto_manage:
+        if console:
+            console.print(
+                "[red]Cannot connect to Ollama.[/red]\n"
+                "Start it manually with [bold]ollama serve[/bold] "
+                "or set [bold]ollama.auto_manage = true[/bold] in your config."
+            )
+        return False
+
+    return start_server(host=host, console=console)
+
+
+def ensure_model(
+    model_name: str,
+    host: str = "http://localhost:11434",
+    timeout: int = 120,
+    console: Console | None = None,
+) -> bool:
+    """Ensure a model is installed, pulling it automatically if missing."""
+    client = OllamaClient(host=host, timeout=timeout)
+
+    # Check if model is already installed
+    try:
+        installed = client.list_models()
+        for m in installed:
+            # Match both exact name and name without tag
+            if m.name == model_name or m.name == f"{model_name}:latest":
+                return True
+            # Also match if user specified without tag
+            if m.name.split(":")[0] == model_name.split(":")[0] and (
+                ":" not in model_name or m.name == model_name
+            ):
+                return True
+    except (ConnectionError, OSError):
+        if console:
+            console.print("[red]Cannot connect to Ollama to check models.[/red]")
+        return False
+
+    # Model not found — pull it
+    if console:
+        console.print(f"[yellow]Model '{model_name}' is not installed.[/yellow]")
+        console.print(f"[dim]Pulling {model_name}...[/dim]")
+
+    try:
+        from rich.progress import (
+            BarColumn,
+            DownloadColumn,
+            Progress,
+            TextColumn,
+            TransferSpeedColumn,
+        )
+
+        if console:
+            with Progress(
+                TextColumn("[progress.description]{task.description}"),
+                BarColumn(),
+                DownloadColumn(),
+                TransferSpeedColumn(),
+                console=console,
+            ) as progress:
+                task = progress.add_task("Downloading", total=None)
+                for update in client.pull(model_name, stream=True):
+                    status = update.get("status", "")
+                    completed = update.get("completed", 0)
+                    total = update.get("total", 0)
+                    if total > 0:
+                        progress.update(task, completed=completed, total=total, description=status)
+                    else:
+                        progress.update(task, description=status)
+                progress.update(task, description="Complete")
+            console.print(f"[green]Model '{model_name}' pulled successfully.[/green]")
+        else:
+            for _update in client.pull(model_name, stream=False):
+                pass
+        return True
+    except Exception as e:
+        if console:
+            console.print(f"[red]Failed to pull '{model_name}': {e}[/red]")
+        return False

mita/plugins/__init__.py

@@ -0,0 +1 @@
+"""MCP plugin system — discover and call tools from MCP servers."""

mita/plugins/client.py

@@ -0,0 +1,152 @@
+"""MCP client wrapper for a single plugin server."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from contextlib import AsyncExitStack
+from typing import Any
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from mita.config.schema import PluginDefinition
+
+
+class MCPPluginClient:
+    """Manages connection to a single MCP server."""
+
+    def __init__(self, plugin: PluginDefinition) -> None:
+        self._plugin = plugin
+        self._session: ClientSession | None = None
+        self._exit_stack: AsyncExitStack | None = None
+
+    @property
+    def name(self) -> str:
+        return self._plugin.name
+
+    @property
+    def transport(self) -> str:
+        return self._plugin.transport
+
+    @property
+    def connected(self) -> bool:
+        return self._session is not None
+
+    async def connect(self, timeout: float = 30.0) -> None:
+        """Connect to the MCP server and initialize the session."""
+        if self._session is not None:
+            return
+
+        self._exit_stack = AsyncExitStack()
+
+        if self._plugin.transport == "stdio":
+            await self._connect_stdio(timeout)
+        elif self._plugin.transport == "sse":
+            await self._connect_sse(timeout)
+        else:
+            raise ValueError(f"Unsupported transport: {self._plugin.transport}")
+
+    async def _connect_stdio(self, timeout: float) -> None:
+        """Connect via stdio transport (subprocess)."""
+        assert self._exit_stack is not None
+
+        if not self._plugin.command:
+            raise ValueError(f"Plugin '{self.name}' requires a command for stdio transport")
+
+        # Build environment: inherit current env, overlay plugin-specific vars
+        env: dict[str, str] = {**os.environ, **self._plugin.env}
+
+        server_params = StdioServerParameters(
+            command=self._plugin.command,
+            args=self._plugin.args,
+            env=env,
+        )
+
+        transport = await self._exit_stack.enter_async_context(stdio_client(server_params))
+        read_stream, write_stream = transport
+
+        session = await self._exit_stack.enter_async_context(
+            ClientSession(read_stream, write_stream)
+        )
+        await asyncio.wait_for(session.initialize(), timeout=timeout)
+        self._session = session
+
+    async def _connect_sse(self, timeout: float) -> None:
+        """Connect via SSE transport (HTTP)."""
+        assert self._exit_stack is not None
+
+        if not self._plugin.url:
+            raise ValueError(f"Plugin '{self.name}' requires a url for SSE transport")
+
+        from mcp.client.sse import sse_client
+
+        transport = await self._exit_stack.enter_async_context(sse_client(self._plugin.url))
+        read_stream, write_stream = transport
+
+        session = await self._exit_stack.enter_async_context(
+            ClientSession(read_stream, write_stream)
+        )
+        await asyncio.wait_for(session.initialize(), timeout=timeout)
+        self._session = session
+
+    async def disconnect(self) -> None:
+        """Disconnect from the MCP server and clean up resources."""
+        if self._exit_stack is not None:
+            await self._exit_stack.aclose()
+            self._exit_stack = None
+        self._session = None
+
+    async def list_tools(self) -> list[dict[str, Any]]:
+        """List tools provided by this MCP server.
+
+        Returns a list of dicts with name, description, and inputSchema.
+        """
+        if self._session is None:
+            raise RuntimeError(f"Plugin '{self.name}' is not connected")
+
+        response = await self._session.list_tools()
+        return [
+            {
+                "name": tool.name,
+                "description": tool.description or "",
+                "inputSchema": tool.inputSchema,
+            }
+            for tool in response.tools
+        ]
+
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any], timeout: float = 120.0
+    ) -> str:
+        """Call a tool on this MCP server.
+
+        Returns the text content from the tool result.
+        """
+        if self._session is None:
+            raise RuntimeError(f"Plugin '{self.name}' is not connected")
+
+        result = await asyncio.wait_for(
+            self._session.call_tool(tool_name, arguments),
+            timeout=timeout,
+        )
+
+        # Extract text content from result
+        parts: list[str] = []
+        for item in result.content:
+            if hasattr(item, "text"):
+                parts.append(item.text)
+            elif hasattr(item, "data"):
+                parts.append(f"[binary data: {getattr(item, 'mimeType', 'unknown')}]")
+            else:
+                parts.append(str(item))
+        return "\n".join(parts)
+
+    async def ping(self, timeout: float = 10.0) -> bool:
+        """Ping the MCP server to check connectivity."""
+        if self._session is None:
+            return False
+        try:
+            await asyncio.wait_for(self._session.send_ping(), timeout=timeout)
+            return True
+        except (TimeoutError, ConnectionError, OSError):
+            return False