Blame for knack.py

#!/usr/bin/env python3
"""CodeSync Knack, standalone: ask the model CodeSync's judgment questions through Ollama.

Standard library only. Needs Ollama (https://ollama.com) running; the model (817 MB) is
pulled on first use. One-liner, nothing to install:

    curl -fsSL https://huggingface.co/devshub-ai/knack/resolve/main/knack.py | python3 - "npm run dev"

The command is only judged, never executed.

Usage:

    python knack.py "npm run dev"                   # the 8 shell-command questions
    python knack.py --turn "fix the failing login test"   # the 21 per-turn questions
    python knack.py --turn "add dark mode" --mode plan --json
    python knack.py serve --port 33727              # CodeSync-compatible POST /v1/systemone

Knack is an Ollama decision model: this script sends all questions of a command or turn in
one ``POST /v1/systemone`` request (needs an Ollama with decision-model support). Without the
script: ``echo "npm run dev" | ollama run devshub/knack --ask serves``.

Env: KNACK_OLLAMA_HOST (default http://127.0.0.1:11434),
     KNACK_MODEL (default devshub/knack; devshub/knack:q4_k_m is smaller, 516 MB;
     hf.co/devshub-ai/knack:Q8_0 pulls the same model from Hugging Face).
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

HOST = os.environ.get("KNACK_OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")
MODEL = os.environ.get("KNACK_MODEL", "devshub/knack")
_READY: set[str] = set()

# question id -> (type, option keys in letter order, what it decides)
QUESTIONS: dict[str, tuple[str, list[str], str]] = {
    # every shell command (state: command, description)
    "risk": ("score", ["0", "1", "2"], "0 read-only, 1 recoverable change, 2 irreversible / remote"),
    "long_running": ("noul", ["yes", "no"], "takes minutes, blocks, polls, watches or serves"),
    "activity": ("choice", ["testing", "building", "installing", "git", "inspecting", "running"],
                 "what kind of work it does"),
    "verifies": ("noul", ["yes", "no"], "a test, type or lint check whose pass proves the change works"),
    "mutates_repo": ("noul", ["yes", "no"], "changes git state (commit, push, merge, rebase, tag)"),
    "deploys": ("noul", ["yes", "no"], "publishes, releases or changes a remote environment"),
    "runtime_check": ("noul", ["yes", "no"], "a health probe of a running service that fails if unhealthy"),
    "serves": ("noul", ["yes", "no"], "keeps running until stopped (server, watcher, tail -f)"),
    # every user turn (state: prompt, current_mode, ...)
    "turn.chat_or_task": ("choice", ["chat", "task"], "conversation or work request"),
    "turn.actionable": ("noul", ["yes", "no"], "asks for something to be done"),
    "turn.question_shaped": ("noul", ["yes", "no"], "phrased as a question"),
    "turn.change_intent": ("choice", ["create", "adjust", "enhance", "fix", "refactor", "none"],
                           "kind of code change asked for"),
    "turn.new_task": ("noul", ["yes", "no"], "starts a new task (vs continuing)"),
    "turn.resume": ("noul", ["yes", "no"], "asks to resume interrupted work"),
    "turn.mission": ("noul", ["yes", "no"], "long multi-step mission"),
    "turn.plan_authoring": ("noul", ["yes", "no"], "asks to write a plan"),
    "turn.screencast": ("choice", ["start", "stop", "none"], "start / stop screen recording"),
    "turn.demo": ("noul", ["yes", "no"], "asks for a demo"),
    "turn.skip_orientation": ("noul", ["yes", "no"], "no codebase exploration needed"),
    "task.intent": ("choice", ["bug_fix", "feature", "refactor", "api_schema", "config_deploy",
                               "test", "documentation", "investigation", "unknown"], "task type"),
    "task.execution_strategy": ("choice", ["direct", "plan", "parallel"], "how to execute"),
    "route.mode": ("choice", ["plan", "build", "ask", "debug"], "best CodeSync mode"),
    "route.tier": ("choice", ["light", "medium", "heavy"], "model size needed"),
    "route.thinking": ("choice", ["off", "low", "medium", "high", "xhigh", "max"],
                       "reasoning effort needed"),
    "route.risk": ("choice", ["low", "medium", "high"], "risk of the requested work"),
    "browser": ("choice", ["yes", "no"], "needs a browser"),
    "messaging": ("choice", ["yes", "no"], "needs to send a message (email, chat)"),
    "delivery.feature_task": ("choice", ["yes", "no"], "a feature delivery task"),
    "delivery.slash_task": ("choice", ["yes", "no"], "a slash-command style task"),
}
BASH = list(QUESTIONS)[:8]  # the shell-command questions above
TURN = list(QUESTIONS)[8:]  # the per-turn questions


def bash_state(command: str, description: str = "") -> dict[str, Any]:
    return {"command": command, "description": description}


def turn_state(prompt: str, mode: str = "build") -> dict[str, Any]:
    return {"prompt": prompt, "current_mode": mode, "recent_activity": "none",
            "conversation_size": "unknown", "open_task": "none", "turns_on_open_task": 0,
            "interrupted_work": "none", "browser_context": False}


def render_state(state: dict[str, Any]) -> str:
    return "\n".join(f"{k}: {v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)}"
                     for k, v in state.items())


def _post(path: str, body: dict[str, Any], timeout: float = 60.0) -> dict[str, Any]:
    req = urllib.request.Request(HOST + path, data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def _spec(qid: str) -> dict[str, Any]:
    """System One question for /v1/systemone (the id is what Knack matches on)."""
    kind, keys, _ = QUESTIONS[qid]
    if kind == "noul":
        return {"type": "noul", "instructions": qid}
    if kind == "score":
        return {"type": "score", "instructions": qid, "criteria": keys}
    return {"type": "choice", "instructions": qid, "criteria": {k: None for k in keys}}


def _answer(qid: str, a: dict[str, Any]) -> dict[str, Any]:
    """{answer, confidence, probabilities, type} from one /v1/systemone answer."""
    kind, keys, _ = QUESTIONS[qid]
    if kind == "noul":
        yes = float(a["noul"])
        probs = {"yes": yes, "no": 1.0 - yes}
    else:
        raw = a.get("probabilities") or {}
        probs = {k: float(raw.get(k, raw.get(str(i), 0.0))) for i, k in enumerate(keys)}
    best = max(probs, key=probs.__getitem__)
    return {"answer": best, "confidence": probs[best], "probabilities": probs, "type": kind}


def ensure_model(model: str = MODEL) -> None:
    """Pull the model into Ollama if it is not there yet (first run only)."""
    if model in _READY:
        return
    _READY.add(model)
    try:
        _post("/api/show", {"model": model}, timeout=10)
        return
    except urllib.error.HTTPError as exc:
        if exc.code != 404:
            raise
    print(f"pulling {model} into Ollama (first run only)...", file=sys.stderr, flush=True)
    _post("/api/pull", {"model": model, "stream": False}, timeout=3600)


def ask(state: dict[str, Any], qids: list[str], model: str = MODEL) -> dict[str, dict[str, Any]]:
    unknown = [q for q in qids if q not in QUESTIONS]
    if unknown:
        raise KeyError(f"Knack was not trained on: {unknown}")
    ensure_model(model)
    try:
        out = _post("/v1/systemone", {"model": model, "state": state,
                                      "questions": {q: _spec(q) for q in qids}})
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="replace")[:300]
        raise RuntimeError(f"Ollama /v1/systemone: HTTP {exc.code} {detail} "
                           "(update Ollama and re-pull the model if this is an old build)") from None
    answers = out.get("answers") or {}
    return {q: _answer(q, answers[q]) for q in qids if q in answers}


def to_systemone(qid: str, a: dict[str, Any]) -> dict[str, Any]:
    """CodeSync / TypeSafe System One answer shape."""
    p, kind = a["probabilities"], a["type"]
    if kind == "noul":
        return {"noul": p["yes"], "confidence": a["confidence"]}
    if kind == "score":
        return {"score": sum(int(k) * v for k, v in p.items()), "confidence": a["confidence"],
                "probabilities": p}
    return {"choice": a["answer"], "confidence": a["confidence"], "probabilities": p}


def serve(port: int, model: str) -> None:
    class H(BaseHTTPRequestHandler):
        def _send(self, code: int, body: dict[str, Any]) -> None:
            data = json.dumps(body).encode()
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def do_GET(self) -> None:  # noqa: N802
            self._send(200, {"ok": True, "model": model, "questions": len(QUESTIONS)})

        def do_POST(self) -> None:  # noqa: N802
            body = json.loads(self.rfile.read(int(self.headers.get("Content-Length") or 0)))
            qids = [q for q in body.get("questions", {}) if q in QUESTIONS]  # others: omitted
            try:
                res = ask(body["state"], qids, model)
            except Exception as exc:  # noqa: BLE001
                return self._send(502, {"error": f"{type(exc).__name__}: {exc}"})
            self._send(200, {"answers": {q: to_systemone(q, a) for q, a in res.items()}})

        def log_message(self, *a: Any) -> None:
            pass

    print(f"Knack ({model} via {HOST}) on http://127.0.0.1:{port}/v1/systemone", flush=True)
    ThreadingHTTPServer(("127.0.0.1", port), H).serve_forever()


def main(argv: list[str] | None = None) -> int:
    argv = sys.argv[1:] if argv is None else argv
    if argv[:1] == ["serve"]:
        sp = argparse.ArgumentParser(prog="knack.py serve")
        sp.add_argument("--port", type=int, default=33727)
        sp.add_argument("--model", default=MODEL)
        a = sp.parse_args(argv[1:])
        serve(a.port, a.model)
        return 0
    ap = argparse.ArgumentParser(description="Ask CodeSync Knack (via Ollama).")
    ap.add_argument("text", help="a shell command, or with --turn a user message")
    ap.add_argument("--turn", action="store_true", help="judge a user message, not a command")
    ap.add_argument("--mode", default="build", help="current CodeSync mode for --turn")
    ap.add_argument("--description", default="", help="what the command is for")
    ap.add_argument("-q", "--question", action="append", help="only these question ids")
    ap.add_argument("--model", default=MODEL)
    ap.add_argument("--json", action="store_true")
    a = ap.parse_args(argv)
    state = turn_state(a.text, a.mode) if a.turn else bash_state(a.text, a.description)
    qids = a.question or (TURN if a.turn else BASH)
    try:
        res = ask(state, qids, a.model)
    except RuntimeError as exc:
        print(exc, file=sys.stderr)
        return 1
    except OSError as exc:
        print(f"cannot reach Ollama at {HOST} ({exc}). Is Ollama running?", file=sys.stderr)
        return 1
    if a.json:
        print(json.dumps(res, indent=1))
        return 0
    w = max(map(len, qids))
    for q, r in res.items():
        print(f"{q:<{w}}  {r['answer']:<13} {r['confidence']:.2f}   {QUESTIONS[q][2]}")
    return 0


if __name__ == "__main__":
    sys.exit(main())