mindxtrain/ui/console.py

"""Talk to what you just trained — Ollama helpers, ported from this house's own consoles.

Provenance: `pythaiml/automindx · ollama_codephreak.py` (the live model picker, streaming `/api/chat`,
a token counter built from the engine's final object, errors written as instructions) and
`Professor-Codephreak/automind · memory.py` (one record per exchange). Two years of interfaces agreed
on one rule and it is kept here: **never block the interface on a model load** — the daemon holds the
weights, the UI asks it questions.

Nothing in this module imports torch. A training box that cannot spare the RAM can still use the
console, because the model lives in Ollama's process, not in this one.
"""
from __future__ import annotations

import json
import os
import time
import urllib.error
import urllib.request
from collections.abc import Iterator
from pathlib import Path
from typing import Any

HOST = os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")

# The standard set, with the defaults Ollama itself uses. The UI renders these; the engine honours
# every one of them, so nothing here is decorative.
DEFAULTS: dict[str, Any] = {
    "temperature": 0.7, "top_p": 0.9, "top_k": 40, "min_p": 0.0,
    "repeat_penalty": 1.1, "repeat_last_n": 64, "num_predict": 256, "num_ctx": 4096,
    "seed": 0, "stop": [], "presence_penalty": 0.0, "frequency_penalty": 0.0,
    "mirostat": 0, "mirostat_tau": 5.0, "mirostat_eta": 0.1,
}

# The imprint gate's decoding. A score measured under anything else is not comparable with the
# numbers in an ascent log, and saying so is cheaper than explaining a mismatch later.
GATE_DECODING: dict[str, Any] = {"temperature": 0.0, "repeat_penalty": 1.3, "top_p": 1.0, "top_k": 0}


def models(host: str = "", timeout: float = 4.0) -> list[str]:
    """Tags the daemon is serving, local first. Empty when it is down — never an exception."""
    try:
        with urllib.request.urlopen((host or HOST).rstrip("/") + "/api/tags", timeout=timeout) as r:
            names = [m["name"] for m in json.loads(r.read().decode()).get("models", [])]
    except Exception:
        return []
    names.sort(key=lambda n: (n.endswith(":cloud"), n))
    return names


def running(host: str = "", timeout: float = 4.0) -> list[dict[str, Any]]:
    """What is resident right now, with its size and expiry — the honest answer to "is it warm?"."""
    try:
        with urllib.request.urlopen((host or HOST).rstrip("/") + "/api/ps", timeout=timeout) as r:
            return json.loads(r.read().decode()).get("models", []) or []
    except Exception:
        return []


def token_line(done: dict[str, Any]) -> str:
    prompt = int(done.get("prompt_eval_count") or 0)
    completion = int(done.get("eval_count") or 0)
    eval_s = (done.get("eval_duration") or 0) / 1e9
    load_s = (done.get("load_duration") or 0) / 1e9
    rate = f" · {completion / eval_s:.2f} tok/s" if eval_s > 0 else ""
    load = f" · {load_s:.1f}s load" if load_s > 0.5 else ""
    return f"{prompt + completion} tokens — {prompt} prompt + {completion} completion{rate}{load}"


def token_stats(done: dict[str, Any]) -> dict[str, Any]:
    prompt = int(done.get("prompt_eval_count") or 0)
    completion = int(done.get("eval_count") or 0)
    eval_s = (done.get("eval_duration") or 0) / 1e9
    return {"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion,
            "tokens_per_s": round(completion / eval_s, 2) if eval_s > 0 else None,
            "load_s": round((done.get("load_duration") or 0) / 1e9, 2),
            "eval_s": round(eval_s, 2), "done_reason": done.get("done_reason")}


def explain(status: int, body: str, model: str, host: str = "") -> str:
    text = (body or "")[:400]
    if status == 404 or "not found" in text.lower():
        return f"**{model}** is not on this daemon:\n\n```bash\nollama pull {model}\n```\n\nOr serve a checkpoint you trained: `mindxtrain serve --to ollama --tag {model}`."
    if status in (401, 402, 403) or "subscription" in text.lower():
        return f"**{model}** looks like a cloud tag needing a subscription. Use a local tag instead."
    if "context" in text.lower():
        return "The prompt is longer than `num_ctx`. Raise it, or send less."
    return f"Ollama answered **{status}** from `{host or HOST}`:\n\n```\n{text}\n```"


def unreachable(host: str = "") -> str:
    return (f"No Ollama at `{host or HOST}`. Start it (`ollama serve`), then either pull a tag or serve "
            "a checkpoint from the **Serve** room. The UI stays up either way — it never holds the weights.")


def chat(messages: list[dict[str, str]], model: str, *, options: dict[str, Any] | None = None,
         keep_alive: str = "10m", host: str = "", timeout: float = 900.0) -> Iterator[tuple[str, dict[str, Any]]]:
    """Stream `(text_so_far, stats)`. `stats` is empty until the final object, which carries the real
    counts. Options are passed through verbatim; the engine implements them, so none are dropped."""
    host = (host or HOST).rstrip("/")
    body = {"model": model, "messages": messages, "stream": True, "keep_alive": keep_alive,
            "options": {k: v for k, v in (options or {}).items() if v not in (None, "", [])}}
    req = urllib.request.Request(f"{host}/api/chat", data=json.dumps(body).encode(), method="POST",
                                 headers={"Content-Type": "application/json"})
    acc = ""
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            for raw in r:
                line = raw.decode("utf-8", errors="replace").strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                except Exception:
                    continue
                acc += (obj.get("message") or {}).get("content", "")
                if obj.get("done"):
                    yield acc, token_stats(obj)
                    return
                if acc:
                    yield acc, {}
    except urllib.error.HTTPError as e:
        yield explain(e.code, e.read().decode("utf-8", errors="replace"), model, host), {"error": True}
        return
    except Exception:
        yield unreachable(host), {"error": True}
        return
    if not acc:
        yield "[the model returned nothing — raise num_predict, or try another tag]", {"error": True}


def record(prompt: str, answer: str, meta: dict[str, Any] | None = None, path: Path | None = None) -> Path | None:
    """One JSONL line per exchange, so a session leaves evidence rather than scrollback."""
    try:
        p = Path(path or os.environ.get("MINDXTRAIN_CONSOLE_LOG") or (Path.home() / ".mindxtrain" / "console.jsonl"))
        p.parent.mkdir(parents=True, exist_ok=True)
        with p.open("a", encoding="utf-8") as fh:
            fh.write(json.dumps({"ts": time.time(), "prompt": prompt[:4000],
                                 "answer": (answer or "")[:8000], **(meta or {})}, ensure_ascii=False) + "\n")
        return p
    except Exception:
        return None