mita-code 0.1.0__py3-none-any.whl

mita/llm/client.py

@@ -0,0 +1,86 @@
+"""LiteLLM client pointing at Ollama for chat completions."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Any
+
+import litellm
+
+from mita.config.schema import MitaConfig
+
+
+class LLMClient:
+    """Thin wrapper around LiteLLM configured for Ollama."""
+
+    def __init__(self, config: MitaConfig) -> None:
+        self._model = f"ollama/{config.model.default}"
+        self._api_base = config.ollama.host
+        self._temperature = config.model.temperature
+        self._max_tokens = config.model.max_tokens
+        self._stream = config.ui.stream
+        self._ollama_options = config.model.ollama_options.to_api_dict()
+
+        # Suppress LiteLLM's verbose logging
+        litellm.suppress_debug_info = True
+
+    @property
+    def model(self) -> str:
+        """The model identifier being used."""
+        return self._model
+
+    async def chat(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> dict[str, Any]:
+        """Send a chat completion request (non-streaming).
+
+        Returns the full response dict from LiteLLM.
+        """
+        kwargs: dict[str, Any] = {
+            "model": self._model,
+            "messages": messages,
+            "temperature": self._temperature,
+            "max_tokens": self._max_tokens,
+            "api_base": self._api_base,
+        }
+        if tools:
+            kwargs["tools"] = tools
+        if self._ollama_options:
+            kwargs["extra_body"] = {"options": self._ollama_options}
+
+        response = await litellm.acompletion(**kwargs)
+        return response  # type: ignore[no-any-return]
+
+    async def stream_chat(
+        self,
+        messages: list[dict[str, Any]],
+        tools: list[dict[str, Any]] | None = None,
+    ) -> AsyncIterator[dict[str, Any]]:
+        """Send a streaming chat completion request.
+
+        Yields delta dicts with content tokens as they arrive.
+        """
+        kwargs: dict[str, Any] = {
+            "model": self._model,
+            "messages": messages,
+            "temperature": self._temperature,
+            "max_tokens": self._max_tokens,
+            "api_base": self._api_base,
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        }
+        if tools:
+            kwargs["tools"] = tools
+        if self._ollama_options:
+            kwargs["extra_body"] = {"options": self._ollama_options}
+
+        response = await litellm.acompletion(**kwargs)
+        async for chunk in response:
+            yield chunk
+
+
+def get_client(config: MitaConfig) -> LLMClient:
+    """Create an LLM client from configuration."""
+    return LLMClient(config)

mita/llm/instructor.py

@@ -0,0 +1,80 @@
+"""Instructor wrapper for structured tool call parsing via JSON mode."""
+
+from __future__ import annotations
+
+from typing import Any, TypeVar
+
+import instructor
+import litellm
+from pydantic import BaseModel
+
+from mita.config.schema import MitaConfig
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class ParsedToolCall(BaseModel):
+    """A tool call parsed from LLM output (llm-local model, no tools dependency)."""
+
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+class ToolCallResponse(BaseModel):
+    """Structured response containing tool calls parsed from LLM output."""
+
+    reasoning: str = ""
+    tool_calls: list[ParsedToolCall] = []
+    text_response: str = ""
+
+
+class InstructorClient:
+    """Instructor-wrapped LiteLLM client for structured output parsing."""
+
+    def __init__(self, config: MitaConfig) -> None:
+        self._model = f"ollama/{config.model.default}"
+        self._api_base = config.ollama.host
+        self._temperature = config.model.temperature
+        self._max_tokens = config.model.max_tokens
+
+        self._client = instructor.from_litellm(
+            litellm.acompletion,
+            mode=instructor.Mode.JSON,
+        )
+
+    async def parse_tool_calls(
+        self,
+        messages: list[dict[str, Any]],
+    ) -> ToolCallResponse:
+        """Parse structured tool calls from LLM output using Instructor JSON mode."""
+        response: ToolCallResponse = await self._client.create(
+            model=self._model,
+            messages=messages,
+            response_model=ToolCallResponse,
+            temperature=self._temperature,
+            max_tokens=self._max_tokens,
+            api_base=self._api_base,
+        )
+        return response
+
+    async def parse_structured(
+        self,
+        messages: list[dict[str, Any]],
+        response_model: type[T],
+    ) -> T:
+        """Parse any structured Pydantic model from LLM output."""
+        response: T = await self._client.create(
+            model=self._model,
+            messages=messages,
+            response_model=response_model,
+            temperature=self._temperature,
+            max_tokens=self._max_tokens,
+            api_base=self._api_base,
+        )
+        return response
+
+
+def get_instructor_client(config: MitaConfig) -> InstructorClient:
+    """Create an Instructor client from configuration."""
+    return InstructorClient(config)

mita/llm/streaming.py

@@ -0,0 +1,58 @@
+"""Token-by-token streaming handler for Rich terminal display."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from mita.llm.client import LLMClient
+
+
+async def stream_to_terminal(
+    client: LLMClient,
+    messages: list[dict[str, Any]],
+    on_token: Callable[[str], None] | None = None,
+    on_complete: Callable[[str], None] | None = None,
+    tools: list[dict[str, Any]] | None = None,
+) -> str:
+    """Stream a chat completion, calling on_token for each chunk.
+
+    Args:
+        client: The LLM client to use.
+        messages: Chat messages.
+        on_token: Callback called with each text token as it arrives.
+        on_complete: Callback called with the full text when streaming is done.
+        tools: Optional tool schemas for the LLM.
+
+    Returns:
+        The complete response text.
+    """
+    full_text = ""
+
+    async for chunk in client.stream_chat(messages, tools=tools):
+        delta = _extract_delta_content(chunk)
+        if delta:
+            full_text += delta
+            if on_token:
+                on_token(delta)
+
+    if on_complete:
+        on_complete(full_text)
+
+    return full_text
+
+
+def _extract_delta_content(chunk: Any) -> str:
+    """Extract text content from a streaming chunk."""
+    try:
+        choices = chunk.choices if hasattr(chunk, "choices") else chunk.get("choices", [])
+        if not choices:
+            return ""
+        delta = choices[0].delta if hasattr(choices[0], "delta") else choices[0].get("delta", {})
+        if hasattr(delta, "content"):
+            return delta.content or ""
+        if isinstance(delta, dict):
+            return delta.get("content", "") or ""
+    except (IndexError, AttributeError, KeyError):
+        pass
+    return ""

mita/memory/__init__.py

@@ -0,0 +1,6 @@
+"""MITA.md memory discovery, loading, and management."""
+
+from mita.memory.discovery import discover_memory_files
+from mita.memory.loader import load_memory
+
+__all__ = ["discover_memory_files", "load_memory"]

mita/memory/discovery.py

@@ -0,0 +1,61 @@
+"""Walk-up-tree discovery of MITA.md memory files."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from mita.config.defaults import PROJECT_CONFIG_DIR
+
+MEMORY_FILENAME = "MITA.md"
+
+
+def get_global_memory_path() -> Path:
+    """Return the path to the global MITA.md file."""
+    return Path.home() / ".config" / "mita" / MEMORY_FILENAME
+
+
+def discover_memory_files(cwd: Path | None = None) -> list[Path]:
+    """Discover MITA.md files by walking up from cwd to the project root.
+
+    Lookup order (returned list, lowest priority first):
+    1. Global: ~/.config/mita/MITA.md
+    2. Project root: <project_root>/MITA.md or <project_root>/.mita/MITA.md
+    3. Directory-level: <cwd>/MITA.md (if different from project root)
+
+    Returns paths ordered from lowest to highest priority (global first).
+    """
+    if cwd is None:
+        cwd = Path.cwd()
+    cwd = cwd.resolve()
+
+    found: list[Path] = []
+
+    # Walk up from cwd, collecting MITA.md files
+    current = cwd
+    while True:
+        candidate = current / MEMORY_FILENAME
+        if candidate.is_file():
+            found.append(candidate)
+
+        # Check .mita/ subdirectory
+        dotmita_candidate = current / PROJECT_CONFIG_DIR / MEMORY_FILENAME
+        if dotmita_candidate.is_file() and dotmita_candidate not in found:
+            found.append(dotmita_candidate)
+
+        # Stop at project root (contains .git or .mita)
+        if (current / ".git").exists() or (current / PROJECT_CONFIG_DIR).exists():
+            break
+
+        parent = current.parent
+        if parent == current:
+            break
+        current = parent
+
+    # Add global memory if it exists and isn't already found
+    global_memory_path = get_global_memory_path()
+    if global_memory_path.is_file() and global_memory_path not in found:
+        found.append(global_memory_path)
+
+    # Reverse so global (lowest priority) is first
+    found.reverse()
+    return found

mita/memory/loader.py

@@ -0,0 +1,76 @@
+"""Read, merge, and truncate MITA.md memory files."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from mita.config.schema import MemorySettings
+from mita.memory.discovery import discover_memory_files
+
+
+def _classify_scope(path: Path) -> str:
+    """Classify a memory file as global, project, or directory scope."""
+    home = Path.home()
+    if str(path).startswith(str(home / ".config" / "mita")):
+        return "global"
+    # If the file is at a project root (parent has .git or .mita), it's project-level
+    parent = path.parent
+    if parent.name == ".mita":
+        parent = parent.parent
+    if (parent / ".git").exists() or (parent / ".mita").exists():
+        return "project"
+    return "directory"
+
+
+def _read_and_truncate(path: Path, max_lines: int) -> tuple[str, bool]:
+    """Read a file and truncate to max_lines. Returns (content, was_truncated)."""
+    lines = path.read_text(encoding="utf-8").splitlines()
+    if len(lines) <= max_lines:
+        return "\n".join(lines), False
+    truncated = lines[:max_lines]
+    truncated.append(f"[... truncated at {max_lines} lines. Keep MITA.md concise.]")
+    return "\n".join(truncated), True
+
+
+def load_memory(
+    cwd: Path | None = None,
+    settings: MemorySettings | None = None,
+) -> str:
+    """Load and merge all discovered MITA.md files into a single string.
+
+    Returns formatted memory content with source annotations, suitable
+    for injection into the system prompt.
+    """
+    if settings is None:
+        settings = MemorySettings()
+
+    files = discover_memory_files(cwd)
+    if not files:
+        return ""
+
+    sections: list[str] = []
+    for path in files:
+        scope = _classify_scope(path)
+        content, _ = _read_and_truncate(path, settings.max_lines_per_file)
+        if content.strip():
+            sections.append(f"<!-- Source: {path} ({scope}) -->\n{content}")
+
+    if not sections:
+        return ""
+
+    return "<memory>\n" + "\n\n".join(sections) + "\n</memory>"
+
+
+def load_memory_raw(cwd: Path | None = None) -> list[tuple[Path, str, str]]:
+    """Load memory files and return as list of (path, scope, content) tuples.
+
+    Useful for the `mita memory show` CLI command.
+    """
+    settings = MemorySettings()
+    files = discover_memory_files(cwd)
+    result: list[tuple[Path, str, str]] = []
+    for path in files:
+        scope = _classify_scope(path)
+        content, _ = _read_and_truncate(path, settings.max_lines_per_file)
+        result.append((path, scope, content))
+    return result

mita/memory/manager.py

@@ -0,0 +1,117 @@
+"""CLI commands for MITA.md memory management."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from pathlib import Path
+
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+
+from mita.config.defaults import _find_project_root
+from mita.memory.discovery import MEMORY_FILENAME, discover_memory_files, get_global_memory_path
+from mita.memory.loader import load_memory_raw
+
+console = Console()
+
+
+def show_memory() -> None:
+    """Display all discovered memory content with source annotations."""
+    entries = load_memory_raw()
+    if not entries:
+        console.print("[dim]No MITA.md files found.[/dim]")
+        return
+
+    for path, scope, content in entries:
+        console.print(
+            Panel(
+                Markdown(content) if content.strip() else "[dim]Empty[/dim]",
+                title=f"[bold]{path}[/bold]",
+                subtitle=f"[dim]{scope}[/dim]",
+                border_style="blue",
+            )
+        )
+
+
+def show_memory_paths() -> None:
+    """Display the paths of all discovered MITA.md files."""
+    files = discover_memory_files()
+    if not files:
+        console.print("[dim]No MITA.md files found.[/dim]")
+        return
+
+    for path in files:
+        console.print(f"  {path}")
+
+
+def edit_memory(scope: str | None = None) -> None:
+    """Open a MITA.md file in $EDITOR.
+
+    Args:
+        scope: "global", "project", or None (nearest/project default).
+    """
+    editor = os.environ.get("EDITOR", "vi")
+
+    if scope == "global":
+        path = get_global_memory_path()
+        path.parent.mkdir(parents=True, exist_ok=True)
+        if not path.exists():
+            path.write_text("# Global Mita Memory\n\n", encoding="utf-8")
+    elif scope == "project":
+        project_root = _find_project_root(Path.cwd())
+        if project_root is None:
+            console.print("[red]Not in a project (no .git or .mita directory found).[/red]")
+            return
+        path = project_root / MEMORY_FILENAME
+        if not path.exists():
+            path.write_text(f"# {project_root.name} — Mita Memory\n\n", encoding="utf-8")
+    else:
+        # Find nearest existing MITA.md, or create at project root
+        files = discover_memory_files()
+        # Filter out global — prefer project/directory level
+        local_files = [f for f in files if "/.config/mita/" not in str(f)]
+        if local_files:
+            path = local_files[-1]  # highest priority (most specific)
+        else:
+            project_root = _find_project_root(Path.cwd())
+            if project_root is None:
+                console.print(
+                    "[red]Not in a project (no .git or .mita directory found). "
+                    "Use --global to edit global memory.[/red]"
+                )
+                return
+            path = project_root / MEMORY_FILENAME
+            if not path.exists():
+                path.write_text(f"# {project_root.name} — Mita Memory\n\n", encoding="utf-8")
+
+    console.print(f"[dim]Opening {path} in {editor}...[/dim]")
+    subprocess.run([editor, str(path)])
+
+
+def add_memory(text: str, scope: str | None = None) -> None:
+    """Append a line of text to a MITA.md file.
+
+    Args:
+        text: Text to append.
+        scope: "global", "project", or None (defaults to project).
+    """
+    if scope == "global":
+        path = get_global_memory_path()
+        path.parent.mkdir(parents=True, exist_ok=True)
+        if not path.exists():
+            path.write_text("# Global Mita Memory\n\n", encoding="utf-8")
+    else:
+        project_root = _find_project_root(Path.cwd())
+        if project_root is None:
+            console.print("[red]Not in a project (no .git or .mita directory found).[/red]")
+            return
+        path = project_root / MEMORY_FILENAME
+        if not path.exists():
+            path.write_text(f"# {project_root.name} — Mita Memory\n\n", encoding="utf-8")
+
+    with open(path, "a", encoding="utf-8") as f:
+        f.write(f"- {text}\n")
+
+    console.print(f"[green]Added to {path}[/green]")

mita/models/__init__.py

@@ -0,0 +1,13 @@
+"""Model management: hardware detection, registry, recommendations, Ollama client."""
+
+from mita.models.hardware import HardwareInfo, detect_hardware
+from mita.models.recommender import ModelRecommendation, recommend_models
+from mita.models.registry import ModelCard
+
+__all__ = [
+    "HardwareInfo",
+    "ModelCard",
+    "ModelRecommendation",
+    "detect_hardware",
+    "recommend_models",
+]