split-stack 0.2.0__py3-none-any.whl

split_stack/ollama_generate.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from split_stack.discovery import discover_models
+from split_stack.local_models import list_local_models
+from split_stack.ollama_errors import format_ollama_error
+from split_stack.routing import route_prompt
+from split_stack.tiering import assign_tiers
+
+
+@dataclass(frozen=True)
+class RouteResult:
+    tier: str
+    model: str
+    ready: bool
+    error: str | None = None
+
+
+@dataclass(frozen=True)
+class AskResult:
+    tier: str
+    model: str
+    text: str
+    ready: bool
+    error: str | None = None
+
+
+def _resolve_model_names(
+    *,
+    base_url: str,
+    model_names: list[str] | None,
+    config_path: str | None = None,
+    only_vram_ok: bool = True,
+) -> tuple[list[str], str | None]:
+    if model_names:
+        return model_names, None
+    resolved, warning = list_local_models(
+        base_url=base_url,
+        config_path=config_path,
+        only_vram_ok=only_vram_ok,
+    )
+    if resolved:
+        return [item.name for item in resolved], warning
+    return discover_models(base_url=base_url), warning
+
+
+def route_prompt_json(
+    prompt: str,
+    *,
+    base_url: str = "http://127.0.0.1:11434",
+    model_names: list[str] | None = None,
+    config_path: str | None = None,
+    only_vram_ok: bool = True,
+    hint: str | None = None,
+) -> RouteResult:
+    try:
+        names, _warning = _resolve_model_names(
+            base_url=base_url,
+            model_names=model_names,
+            config_path=config_path,
+            only_vram_ok=only_vram_ok,
+        )
+        if not names:
+            return RouteResult(tier="", model="", ready=False, error="No models available")
+        tiers = assign_tiers(names)
+        tier, model = route_prompt(prompt, tiers, hint=hint)
+        return RouteResult(tier=tier.value, model=model, ready=True)
+    except Exception as exc:
+        return RouteResult(tier="", model="", ready=False, error=str(exc))
+
+
+def generate_text(
+    model: str,
+    prompt: str,
+    *,
+    base_url: str = "http://127.0.0.1:11434",
+    timeout_seconds: int = 60,
+) -> str:
+    try:
+        import requests
+    except ImportError as exc:
+        raise RuntimeError(
+            "generate_text requires optional dependency: pip install split-stack[ollama]"
+        ) from exc
+
+    try:
+        response = requests.post(
+            f"{base_url.rstrip('/')}/api/generate",
+            json={"model": model, "prompt": prompt, "stream": False},
+            timeout=timeout_seconds,
+        )
+        response.raise_for_status()
+    except Exception as exc:
+        raise RuntimeError(format_ollama_error(exc, model=model, base_url=base_url)) from exc
+    payload = response.json() or {}
+    return payload.get("response", "").strip()
+
+
+def ask_prompt_json(
+    prompt: str,
+    *,
+    base_url: str = "http://127.0.0.1:11434",
+    model_names: list[str] | None = None,
+    timeout_seconds: int = 60,
+    config_path: str | None = None,
+    only_vram_ok: bool = True,
+    hint: str | None = None,
+) -> AskResult:
+    routed = route_prompt_json(
+        prompt,
+        base_url=base_url,
+        model_names=model_names,
+        config_path=config_path,
+        only_vram_ok=only_vram_ok,
+        hint=hint,
+    )
+    if not routed.ready:
+        return AskResult(tier="", model="", text="", ready=False, error=routed.error)
+    try:
+        text = generate_text(
+            routed.model,
+            prompt,
+            base_url=base_url,
+            timeout_seconds=timeout_seconds,
+        )
+        return AskResult(tier=routed.tier, model=routed.model, text=text, ready=True)
+    except Exception as exc:
+        return AskResult(
+            tier=routed.tier,
+            model=routed.model,
+            text="",
+            ready=False,
+            error=str(exc),
+        )

split_stack/poc_models.py

@@ -0,0 +1,131 @@
+"""Default model stacks for POC demos and compare benchmarks."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from split_stack.community_picks import focus_stack, vram_tier_for_profile
+from split_stack.discovery import list_model_inventory
+from split_stack.presets import recommended_models
+
+DEFAULT_POC_STACK = ["gemma4:e4b", "qwen3:8b", "qwen3:14b"]
+QWEN_ONLY_STACK = ["qwen3:4b", "qwen3:8b", "qwen3:14b"]
+
+
+@dataclass(frozen=True)
+class StackPreset:
+    id: str
+    label: str
+    models: tuple[str, ...]
+    description: str
+
+
+STACK_PRESETS: tuple[StackPreset, ...] = (
+    StackPreset(
+        id="mixed_12gb",
+        label="Mixed 12 GB (Gemma + Qwen)",
+        models=tuple(DEFAULT_POC_STACK),
+        description="Gemma lookup, Qwen 8B medium, Qwen 14B complex",
+    ),
+    StackPreset(
+        id="qwen_only",
+        label="Qwen only (4B / 8B / 14B)",
+        models=tuple(QWEN_ONLY_STACK),
+        description="Single-family ladder",
+    ),
+    StackPreset(
+        id="recommended_12gb",
+        label="Full 12 GB specialist",
+        models=tuple(recommended_models("workstation_12gb")),
+        description="Gemma + Qwen + DeepSeek R1 for reasoning",
+    ),
+    StackPreset(
+        id="community_agentic",
+        label="Reddit agentic (M tier)",
+        models=tuple(),
+        description="r/LocalLLaMA Apr 2026 — Gemma lookup + Qwen spine for agent loops",
+    ),
+    StackPreset(
+        id="from_inventory",
+        label="From your Ollama (auto ladder)",
+        models=tuple(),
+        description="Picks small/mid/large tags from API + disk manifests",
+    ),
+)
+
+
+def list_stack_presets() -> tuple[StackPreset, ...]:
+    return STACK_PRESETS
+
+
+def models_for_preset(
+    preset_id: str,
+    *,
+    base_url: str = "http://127.0.0.1:11434",
+    profile: str = "workstation_12gb",
+) -> list[str]:
+    if preset_id == "from_inventory":
+        inventory = list_model_inventory(base_url=base_url)
+        if inventory.suggested_stack:
+            return list(inventory.suggested_stack)
+        return list(DEFAULT_POC_STACK)
+    if preset_id == "community_agentic":
+        tier = vram_tier_for_profile(profile)
+        stack = focus_stack("agentic", vram_tier=tier)
+        if stack and stack.models:
+            return list(stack.models)
+        return list(DEFAULT_POC_STACK)
+    for item in STACK_PRESETS:
+        if item.id == preset_id:
+            return list(item.models)
+    valid = ", ".join(item.id for item in STACK_PRESETS)
+    raise ValueError(f"Unknown stack preset '{preset_id}'. Valid: {valid}")
+
+
+def available_model_pool(
+    *,
+    base_url: str = "http://127.0.0.1:11434",
+    source: str = "both",
+) -> tuple[list[str], str | None]:
+    """Return model names from Ollama API, disk manifests, or both."""
+    inventory = list_model_inventory(base_url=base_url)
+    if source == "api":
+        pool = list(inventory.api_models)
+    elif source == "disk":
+        pool = list(inventory.disk_models)
+    else:
+        pool = sorted(set(inventory.api_models) | set(inventory.disk_models))
+    return pool, inventory.note
+
+
+def resolve_installed_stack(
+    installed: list[str],
+    *,
+    preset_id: str = "mixed_12gb",
+    base_url: str = "http://127.0.0.1:11434",
+) -> tuple[list[str], str | None]:
+    """Pick preset models that exist in the installed pool; warn when falling back."""
+    desired = models_for_preset(preset_id, base_url=base_url)
+    installed_set = set(installed)
+    matched = [name for name in desired if name in installed_set]
+    if len(matched) >= 2:
+        return matched, None
+
+    if installed:
+        from split_stack.model_registry import load_registry, model_weight
+
+        registry = load_registry()
+        ranked = sorted(installed, key=lambda name: model_weight(name, registry))
+        if len(ranked) >= 2:
+            warning = (
+                f"Preset '{preset_id}' not fully available ({', '.join(desired)}). "
+                f"Using: {', '.join(ranked)}"
+            )
+            return ranked, warning
+        warning = (
+            f"Preset '{preset_id}' not fully available. "
+            f"Using only {ranked[0]} — need 2+ models for compare spread."
+        )
+        return ranked, warning
+
+    return desired, f"Using preset list (not verified): {', '.join(desired)}"

split_stack/presets.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from split_stack.model_registry import normalize_deployment_profile
+from split_stack.models import TierMap
+from split_stack.tiering import assign_tiers
+
+
+@dataclass(frozen=True)
+class RecommendedStack:
+    profile: str
+    models: tuple[str, ...]
+    description: str
+
+
+RECOMMENDED_STACKS: dict[str, RecommendedStack] = {
+    "workstation_8gb": RecommendedStack(
+        profile="workstation_8gb",
+        models=("gemma4:e4b", "qwen3:8b"),
+        description="Gemma min + Qwen 8b max (flat but honest on 8 GB)",
+    ),
+    "workstation_12gb": RecommendedStack(
+        profile="workstation_12gb",
+        models=("gemma4:e4b", "qwen3:8b", "qwen3:14b", "deepseek-r1:8b"),
+        description="Gemma lookup, Qwen mid/complex, DeepSeek R1 reasoning",
+    ),
+    "workstation_16gb": RecommendedStack(
+        profile="workstation_16gb",
+        models=("gemma4:e4b", "qwen3:8b", "qwen3:14b"),
+        description="Gemma lookup + Qwen mid/complex; add coder/reason tags via models=",
+    ),
+    "workstation_24gb": RecommendedStack(
+        profile="workstation_24gb",
+        models=(
+            "gemma4:e4b",
+            "qwen3:8b",
+            "qwen3:14b",
+            "qwen3:30b-a3b",
+            "deepseek-coder:6.7b",
+        ),
+        description="Full mixed ladder with MoE top and code specialist",
+    ),
+    "workstation_32gb": RecommendedStack(
+        profile="workstation_32gb",
+        models=(
+            "gemma4:e4b",
+            "qwen3:8b",
+            "qwen3:14b",
+            "qwen3:30b-a3b",
+            "deepseek-coder:6.7b",
+            "deepseek-r1:8b",
+        ),
+        description="5090 class: MoE + separate reasoning and code specialists",
+    ),
+}
+
+
+def list_recommended_stacks() -> tuple[RecommendedStack, ...]:
+    return tuple(RECOMMENDED_STACKS[name] for name in sorted(RECOMMENDED_STACKS))
+
+
+def recommended_models(profile: str, *, quant: str | None = None) -> list[str]:
+    profile_name = normalize_deployment_profile(profile)
+    stack = RECOMMENDED_STACKS.get(profile_name)
+    if stack is None:
+        valid = ", ".join(sorted(RECOMMENDED_STACKS))
+        raise ValueError(f"Unknown profile '{profile}'. Valid workstation stacks: {valid}")
+    from split_stack.quantization import expand_models_for_quant
+
+    return expand_models_for_quant(list(stack.models), profile_name, quant)
+
+
+def assign_recommended_tiers(profile: str) -> TierMap:
+    return assign_tiers(recommended_models(profile))

split_stack/quantization.py

@@ -0,0 +1,137 @@
+"""Quantization assumptions for VRAM sizing — not per-prompt routing.
+
+Ollama tags (``gemma4:e4b``) do not encode quant. ``quant=`` tells split-stack which
+pull format you use so VRAM filters and QAT-aware stack suggestions stay honest.
+"""
+
+from __future__ import annotations
+
+import os
+
+QUANT_MODES: tuple[str, ...] = ("default", "qat", "qat_mobile", "bf16")
+
+# Runtime memory (GB) from Unsloth Gemma 4 QAT docs — UD-Q4_K_XL, not naive Q4_0.
+_GEMMA4_QAT_RUNTIME_GB: dict[str, int] = {
+    "gemma4:e2b": 3,
+    ":e2b": 3,
+    "gemma4:e4b": 5,
+    ":e4b": 5,
+    "gemma4:12b": 7,
+    "gemma4:26b-a4b": 15,
+    "gemma4:26b": 15,
+    "gemma4:31b": 18,
+}
+
+# Google mobile mixture QAT (UD-Q2_K_XL class).
+_GEMMA4_QAT_MOBILE_RUNTIME_GB: dict[str, int] = {
+    "gemma4:e2b": 3,
+    ":e2b": 3,
+    "gemma4:e4b": 4,
+    ":e4b": 4,
+}
+
+# BF16 original sizes (Unsloth table, rounded up for filter headroom).
+_GEMMA4_BF16_RUNTIME_GB: dict[str, int] = {
+    "gemma4:e2b": 10,
+    ":e2b": 10,
+    "gemma4:e4b": 16,
+    ":e4b": 16,
+    "gemma4:12b": 24,
+    "gemma4:26b-a4b": 51,
+    "gemma4:26b": 51,
+    "gemma4:31b": 62,
+}
+
+# Extra models that fit when Gemma pulls use QAT int4 (Unsloth hardware table).
+QAT_STACK_ADDITIONS: dict[str, tuple[str, ...]] = {
+    "workstation_16gb": ("gemma4:26b-a4b",),
+    "workstation_24gb": ("gemma4:31b",),
+}
+
+
+def normalize_quant_mode(name: str | None) -> str:
+    if not name:
+        return "default"
+    lowered = name.strip().lower().replace("-", "_")
+    aliases = {
+        "off": "default",
+        "none": "default",
+        "q4": "qat",
+        "q4_qat": "qat",
+        "qat_q4": "qat",
+        "mobile": "qat_mobile",
+        "qat_mobile_mixture": "qat_mobile",
+        "fp16": "bf16",
+        "full": "bf16",
+    }
+    lowered = aliases.get(lowered, lowered)
+    if lowered not in QUANT_MODES:
+        valid = ", ".join(QUANT_MODES)
+        raise ValueError(f"Unknown quant mode '{name}'. Valid modes: {valid}")
+    return lowered
+
+
+def quant_from_env() -> str | None:
+    raw = os.environ.get("SPLIT_STACK_QUANT", "").strip()
+    return raw or None
+
+
+def _lookup_table_vram(name: str, table: dict[str, int]) -> int | None:
+    lowered = name.lower()
+    best_gb: int | None = None
+    best_len = -1
+    for key, gb in table.items():
+        if key in lowered and len(key) > best_len:
+            best_gb = gb
+            best_len = len(key)
+    return best_gb
+
+
+def adjust_vram_for_quant(
+    name: str,
+    base_vram_gb: int | None,
+    quant_mode: str | None,
+) -> int | None:
+    """Return effective VRAM for feasibility checks; does not change routing weight."""
+    mode = normalize_quant_mode(quant_mode)
+    if mode == "default":
+        return base_vram_gb
+    tables = {
+        "qat": _GEMMA4_QAT_RUNTIME_GB,
+        "qat_mobile": _GEMMA4_QAT_MOBILE_RUNTIME_GB,
+        "bf16": _GEMMA4_BF16_RUNTIME_GB,
+    }
+    matched = _lookup_table_vram(name, tables[mode])
+    if matched is not None:
+        return matched
+    return base_vram_gb
+
+
+def expand_models_for_quant(models: list[str], profile: str, quant_mode: str | None) -> list[str]:
+    """Add QAT-feasible models to a recommended stack (Gemma 4 only today)."""
+    if normalize_quant_mode(quant_mode) != "qat":
+        return models
+    profile_name = profile.strip().lower()
+    extras = QAT_STACK_ADDITIONS.get(profile_name, ())
+    out = list(models)
+    for name in extras:
+        if name not in out:
+            out.append(name)
+    return out
+
+
+def pull_guidance_lines(quant_mode: str | None) -> list[str]:
+    """Short pull hints for docs/CLI — not import-time spam."""
+    mode = normalize_quant_mode(quant_mode)
+    if mode == "default":
+        return []
+    if mode == "qat":
+        return [
+            "Gemma 4 QAT: prefer Unsloth UD-Q4_K_XL GGUF over naive Google Q4_0 for llama.cpp/Ollama imports.",
+            "Collections: google/gemma-4-qat-q4_0, unsloth/gemma-4-qat — see docs/LOCAL_MODELS.md",
+        ]
+    if mode == "qat_mobile":
+        return [
+            "Gemma 4 mobile mixture QAT: google/gemma-4-qat-mobile (UD-Q2_K_XL on E2B/E4B).",
+        ]
+    return ["Gemma 4 BF16 pulls need datacenter profile or custom vram_gb in config."]