caudate-cli 0.1.0__py3-none-any.whl

llm/models.py

@@ -0,0 +1,320 @@
+"""Model registry — detect Ollama models and tag their capabilities.
+
+Cognos runs local-first on Ollama but LiteLLM lets us also talk to cloud
+models. The registry exposes:
+  - available models (from the Ollama API, plus any cloud models known)
+  - per-model capability flags (tool calling, JSON mode, context window)
+  - presets: `fast`, `balanced`, `powerful` — resolved against what's installed
+
+Callers pass a preset name (e.g. `--model fast`) or a concrete model ID
+(e.g. `ollama/gemma3:27b`). The registry resolves presets but passes concrete
+IDs through untouched.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+import subprocess
+from dataclasses import dataclass, field
+from typing import Literal
+
+logger = logging.getLogger(__name__)
+
+
+# ──────────────────────── VRAM estimation ─────────────────────────────
+#
+# Ported from LocalForge's lib/models/vram-estimates.ts. Conservative
+# Q4-quantization figures with a small overhead for context. Used by:
+#   - `cognos models` CLI to flag ✅ / ⚠️ / ❌ next to each entry
+#   - `cognos forge` to refuse to start orchestration when no installed
+#     model fits the local GPU
+#   - the model selector preset logic below
+
+_SIZE_TO_VRAM_MB: list[tuple[float, float, int]] = [
+    (0,    2,    1500),
+    (2,    4,    2500),
+    (4,    9,    5500),
+    (9,    13,   8500),
+    (13,   16,   10500),
+    (16,   24,   15000),
+    (24,   35,   21000),
+    (35,   50,   30000),
+    (50,   80,   44000),
+    (80,   200,  100000),
+    (200,  500,  240000),
+]
+
+# Boundary-anchored param-count match (e.g. "-31b", ":7b", "_8b", " 70b").
+# Same regex as the TS port.
+_PARAM_RE = re.compile(
+    r"(?:^|[^a-zA-Z0-9])(\d+(?:\.\d+)?)b(?:[^a-zA-Z0-9]|$)",
+    re.IGNORECASE,
+)
+
+
+@dataclass
+class VramEstimate:
+    params_b: float
+    vram_mb: int
+
+
+FitStatus = Literal["fits", "tight", "wont-fit"]
+
+
+def estimate_vram(model_id: str) -> VramEstimate | None:
+    """Parse a model id like ``ollama/qwen2.5:14b`` and return its
+    estimated VRAM requirement at Q4. ``None`` if no parameter count
+    can be parsed (cloud-only models, embedding models, etc.)."""
+    if not model_id:
+        return None
+    m = _PARAM_RE.search(model_id)
+    if not m:
+        return None
+    try:
+        params_b = float(m.group(1))
+    except ValueError:
+        return None
+    if params_b <= 0:
+        return None
+    for (lo, hi, mb) in _SIZE_TO_VRAM_MB:
+        if lo < params_b <= hi:
+            return VramEstimate(params_b=params_b, vram_mb=mb)
+    return None
+
+
+def compare_to_available(estimate_mb: int, available_mb: int) -> FitStatus:
+    """≤ 70 % = ``fits``; 70–100 % = ``tight``; over = ``wont-fit``."""
+    if available_mb <= 0:
+        return "wont-fit"
+    ratio = estimate_mb / available_mb
+    if ratio <= 0.7:
+        return "fits"
+    if ratio <= 1.0:
+        return "tight"
+    return "wont-fit"
+
+
+def pick_best_fit(
+    model_ids: list[str], available_mb: int,
+) -> tuple[str, VramEstimate, FitStatus] | None:
+    """Largest comfortable fit (≤ 70 %); fall back to largest tight
+    fit if none. Never returns a model that won't fit. Cloud-only ids
+    that don't carry a param count are skipped."""
+    candidates: list[tuple[str, VramEstimate, FitStatus]] = []
+    for mid in model_ids:
+        est = estimate_vram(mid)
+        if est is None:
+            continue
+        candidates.append((mid, est, compare_to_available(est.vram_mb, available_mb)))
+    if not candidates:
+        return None
+    fitting = sorted(
+        [c for c in candidates if c[2] == "fits"],
+        key=lambda c: c[1].params_b, reverse=True,
+    )
+    if fitting:
+        return fitting[0]
+    tight = sorted(
+        [c for c in candidates if c[2] == "tight"],
+        key=lambda c: c[1].params_b, reverse=True,
+    )
+    return tight[0] if tight else None
+
+
+def detect_available_vram_mb() -> int | None:
+    """Best-effort detection via ``nvidia-smi``. Returns total VRAM of
+    GPU 0 in MB, or ``None`` if unavailable (CPU-only host, AMD, etc.)."""
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=memory.total",
+             "--format=csv,noheader,nounits"],
+            stderr=subprocess.DEVNULL, timeout=2,
+        ).decode().strip().splitlines()
+        if not out:
+            return None
+        return int(float(out[0].strip()))
+    except (FileNotFoundError, subprocess.SubprocessError, ValueError):
+        return None
+
+
+# Glyphs for CLI fit indicators
+FIT_GLYPH = {"fits": "OK", "tight": "TIGHT", "wont-fit": "NO-FIT"}
+
+
+@dataclass
+class ModelInfo:
+    id: str  # LiteLLM-compatible model id, e.g. "ollama/gemma3:27b"
+    name: str  # bare name, e.g. "gemma3:27b"
+    provider: str  # "ollama" | "anthropic" | "openai" | …
+    supports_tool_calling: bool = False
+    supports_json_mode: bool = False
+    context_window: int = 8192
+    size_bytes: int = 0
+    tags: list[str] = field(default_factory=list)
+
+
+# Heuristic capability flags, based on model name substrings.
+# Conservative defaults — override per-id as needed.
+_NATIVE_TOOL_CALLING = ("qwen3-coder", "qwen2.5-coder", "llama3.1", "llama3.2",
+                        "llama3", "mistral", "nemotron", "kimi", "glm", "qwen3-vl")
+_NO_NATIVE_TOOL_CALLING = ("gemma", "gemma2", "gemma3", "gemma4")
+_LARGE_CONTEXT = ("qwen3", "kimi", "nemotron", "gemma3", "gemma4")
+
+
+def _classify(name: str) -> tuple[bool, bool, int]:
+    """Return (supports_tool_calling, supports_json_mode, context_window)."""
+    lower = name.lower()
+    tool_calling = True
+    for bad in _NO_NATIVE_TOOL_CALLING:
+        if bad in lower:
+            tool_calling = False
+            break
+    if tool_calling:
+        if not any(good in lower for good in _NATIVE_TOOL_CALLING):
+            # Unknown model — be conservative
+            tool_calling = False
+
+    json_mode = True  # Ollama supports JSON output via response_format for most models
+    context = 32000 if any(x in lower for x in _LARGE_CONTEXT) else 8192
+    return tool_calling, json_mode, context
+
+
+class ModelRegistry:
+    """Detect local + cloud models and expose capability info."""
+
+    def __init__(self):
+        self._models: dict[str, ModelInfo] = {}
+
+    async def refresh(self) -> None:
+        """Re-query Ollama for its installed models."""
+        self._models.clear()
+        for m in await _ollama_list():
+            self._models[m.id] = m
+        for m in _known_cloud_models():
+            self._models[m.id] = m
+
+    def models(self) -> list[ModelInfo]:
+        return list(self._models.values())
+
+    def get(self, model_id: str) -> ModelInfo | None:
+        return self._models.get(model_id)
+
+    def resolve(self, name_or_preset: str) -> str:
+        """Resolve a preset name or pass through a concrete model id."""
+        preset = name_or_preset.lower()
+        if preset == "fast":
+            return self._pick_smallest_ollama() or name_or_preset
+        if preset == "balanced":
+            return self._pick_balanced_ollama() or name_or_preset
+        if preset == "powerful":
+            return self._pick_largest_ollama() or name_or_preset
+        return name_or_preset
+
+    # ------------------------------------------------------------------
+
+    def _pick_smallest_ollama(self) -> str | None:
+        candidates = [m for m in self._models.values() if m.provider == "ollama"]
+        if not candidates:
+            return None
+        return min(candidates, key=lambda m: m.size_bytes or float("inf")).id
+
+    def _pick_largest_ollama(self) -> str | None:
+        candidates = [m for m in self._models.values() if m.provider == "ollama"]
+        if not candidates:
+            return None
+        return max(candidates, key=lambda m: m.size_bytes or 0).id
+
+    def _pick_balanced_ollama(self) -> str | None:
+        # Prefer a mid-size model with tool calling if available
+        candidates = sorted(
+            [m for m in self._models.values() if m.provider == "ollama"],
+            key=lambda m: m.size_bytes or 0,
+        )
+        if not candidates:
+            return None
+        tooled = [m for m in candidates if m.supports_tool_calling]
+        if tooled:
+            return tooled[len(tooled) // 2].id
+        return candidates[len(candidates) // 2].id
+
+
+# --- Ollama API via CLI (no extra deps) ---
+
+
+async def _ollama_list() -> list[ModelInfo]:
+    """Query `ollama list` and return parsed ModelInfo entries."""
+    try:
+        proc = await asyncio.create_subprocess_shell(
+            "ollama list",
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=10)
+    except Exception as e:
+        logger.debug(f"ollama list failed: {e}")
+        return []
+
+    if proc.returncode != 0:
+        return []
+
+    models: list[ModelInfo] = []
+    lines = stdout.decode().splitlines()
+    if len(lines) < 2:
+        return []
+    # Skip header row
+    for line in lines[1:]:
+        parts = line.split()
+        if len(parts) < 3:
+            continue
+        name = parts[0]
+        # size is usually parts[2] + parts[3] (e.g. "17 GB")
+        size_str = parts[2] + " " + parts[3] if len(parts) >= 4 else "0 MB"
+        size_bytes = _parse_size(size_str)
+
+        tool_calling, json_mode, context = _classify(name)
+        models.append(ModelInfo(
+            id=f"ollama/{name}",
+            name=name,
+            provider="ollama",
+            supports_tool_calling=tool_calling,
+            supports_json_mode=json_mode,
+            context_window=context,
+            size_bytes=size_bytes,
+        ))
+    return models
+
+
+def _parse_size(s: str) -> int:
+    """Parse '17 GB' / '986 MB' into bytes."""
+    try:
+        num_str, unit = s.split()
+        num = float(num_str)
+    except (ValueError, AttributeError):
+        return 0
+    unit = unit.upper()
+    mult = {"B": 1, "KB": 1024, "MB": 1024**2, "GB": 1024**3, "TB": 1024**4}.get(unit, 1)
+    return int(num * mult)
+
+
+def _known_cloud_models() -> list[ModelInfo]:
+    """Hard-coded entries for common cloud models so callers can pick them."""
+    return [
+        ModelInfo(
+            id="claude-opus-4-7",
+            name="claude-opus-4-7",
+            provider="anthropic",
+            supports_tool_calling=True,
+            supports_json_mode=True,
+            context_window=200000,
+        ),
+        ModelInfo(
+            id="claude-sonnet-4-6",
+            name="claude-sonnet-4-6",
+            provider="anthropic",
+            supports_tool_calling=True,
+            supports_json_mode=True,
+            context_window=200000,
+        ),
+    ]