CoolFace
Apppublic

build-small-hackathon/microfactory-lab

sourceHugging Facemitupdated 3mo agoView on Hugging Face
2likes
llm.py162 linesDownload Raw Back to core
1"""Local Ollama client (real calls, not mocks).2 3Fully local Gemma inference → earns Off the Grid + Llama Champion by4construction (Ollama runs on llama.cpp). The UI surfaces which model ran and5whether it was a real call or the deterministic fallback (never crash the demo).6 7`gemma4:e4b` is the default (= gemma4:latest, 9.6GB). `gemma4:e2b` if CPU8latency demands. NOTE: `gemma4:4b` does NOT exist — never use that tag.9"""10 11from __future__ import annotations12 13import json14import os15 16MODEL = os.environ.get("CHIEF_ENGINEER_MODEL", "gemma4:e4b")17MODAL_API_URL = os.environ.get("CHIEF_ENGINEER_MODAL_URL",18    "https://kylebrodeur--microfactory-node-inference-serve.modal.run/v1/chat/completions")19 20# Backend select. Default "ollama" keeps local/recording behavior IDENTICAL.21# CHIEF_ENGINEER_BACKEND is the *initial default* (e.g. a Space var = zerogpu); it is22# NOT a hard lock. The in-app model switcher changes the backend at runtime by setting23# the env var, so routing reads it dynamically via _backend() — a fixed Space var would24# otherwise freeze the ZeroGPU<->Modal switch. BACKEND keeps the startup value for25# back-compat/logging. Unknown/import-fail → ollama.26BACKEND = os.environ.get("CHIEF_ENGINEER_BACKEND", "ollama").lower()27 28 29def _backend() -> str:30    """Active backend, read dynamically so the model switcher's selection takes effect31    at runtime (the dropdown sets CHIEF_ENGINEER_BACKEND via app._apply_model_choice)."""32    return os.environ.get("CHIEF_ENGINEER_BACKEND", BACKEND).lower()33 34 35try:36    import ollama  # type: ignore37except Exception:  # pragma: no cover38    ollama = None  # type: ignore39 40 41def _zerogpu():42    """Lazily import the ZeroGPU backend; None unless selected and importable."""43    if _backend() != "zerogpu":44        return None45    try:46        from . import llm_zerogpu  # heavy deps are import-guarded inside47        return llm_zerogpu48    except Exception:49        return None50 51 52def _modal_api():53    """Lazily check if Modal API backend is selected."""54    if _backend() != "modal":55        return None56    return True  # Modal API is always available (HTTP endpoint)57 58 59def _forced_offline() -> bool:60    """Force the deterministic fallback path regardless of any daemon/backend.61    Read dynamically (not cached) so tests can toggle it. Used by the offline62    core suite so `make test` never touches Ollama, even when `ollama serve` is up."""63    return os.environ.get("CHIEF_ENGINEER_OFFLINE", "").lower() in ("1", "true", "yes")64 65 66def is_available() -> bool:67    """True if the active backend can serve a real call."""68    if _forced_offline():69        return False70    zg = _zerogpu()71    if zg is not None:72        return zg.is_available()73    if _modal_api():74        return True75    if ollama is None:76        return False77    try:78        ollama.list()79        return True80    except Exception:81        return False82 83 84def backend_status() -> str:85    zg = _zerogpu()86    if zg is not None:87        return zg.backend_status()88    if _modal_api():89        return f"<span style='color:var(--ao-green);'>●</span> live · Modal API (remote GPU)"90    return (f"<span style='color:var(--ao-green);'>●</span> live · {MODEL} (local Ollama)"91            if is_available() else92            f"<span style='color:var(--ao-yellow);'>●</span> offline fallback · "93            f"{MODEL} unreachable (deterministic)")94 95 96def warm_up() -> str:97    """Pay the model's cold start now (off-camera), so the first real BUILD is fast.98    On ZeroGPU this enters the GPU window and loads the model; on Ollama/fallback it is99    a cheap no-op. On Modal API it's a no-op (Modal handles its own warm-up).100    Returns the (post-load) backend status. Never raises."""101    zg = _zerogpu()102    if zg is not None:103        try:104            return zg.warm()105        except Exception:106            return backend_status()107    return backend_status()108 109 110def chat_json(system: str, user: str, temperature: float = 0.4) -> dict | None:111    """One JSON-mode chat turn. Returns parsed dict, or None to signal fallback."""112    if _forced_offline():113        return None114    zg = _zerogpu()115    if zg is not None:116        try:117            return zg.chat_json(system, user, temperature)118        except Exception:119            return None120    if _modal_api():121        try:122            import urllib.request123            body = json.dumps({124                "messages": [{"role": "user", "content": f"{system}\n\n{user}"}],125                "max_tokens": 512,126                "temperature": temperature,127            }).encode()128            req = urllib.request.Request(MODAL_API_URL, data=body,129                headers={"Content-Type": "application/json"})130            with urllib.request.urlopen(req, timeout=120) as resp:131                data = json.loads(resp.read())132                text = data["choices"][0]["message"]["content"].strip()133                if text.startswith("```"):134                    text = text.strip("`").lstrip()135                    if text[:4].lower() == "json":136                        text = text[4:]137                return json.loads(text)138        except Exception:139            return None140    if not is_available():141        return None142    try:143        resp = ollama.chat(144            model=MODEL,145            messages=[146                {"role": "system", "content": system},147                {"role": "user", "content": user},148            ],149            format="json",150            options={"temperature": temperature},151        )152        content = resp["message"]["content"].strip()153        # Fence-strip safety net (GEMMA-STEERING Technique 2): small Gemmas can154        # wrap JSON in ```json fences even in JSON mode. Strip before parsing.155        if content.startswith("```"):156            content = content.strip("`").lstrip()157            if content[:4].lower() == "json":158                content = content[4:]159        return json.loads(content)160    except Exception:161        return None162