mu-agent 0.1.0

package/kyrn/local-judge/requirements.lock.txt

@@ -0,0 +1,26 @@
+anyio==4.15.1
+attrs==26.1.0
+cattrs==26.2.0
+certifi==2026.7.22
+click==8.5.0
+coremltools==9.0
+filelock==4.0.1
+fsspec==2026.9.0
+h11==0.16.0
+hf-xet==1.6.0
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==1.32.0
+idna==3.20
+laya-coreml==0.1.0
+mpmath==1.3.0
+numpy==2.1.3
+packaging==26.3
+protobuf==7.36.2
+pyaml==26.7.0
+pyyaml==6.0.3
+safetensors==0.8.0
+sympy==1.14.0
+tokenizers==0.23.2
+tqdm==4.70.1
+typing-extensions==4.16.0

package/kyrn/local-judge/server.py

@@ -0,0 +1,332 @@
+#!/usr/bin/env python3
+"""mu local judge sidecar.
+
+Serves a Laya Core ML checkpoint (https://github.com/mizorewww/laya-coreml) over
+the same wire contract as the AI Gateway `evaluation-model` endpoint, so the
+mu kernel can swap Jev for a local model without translating anything:
+
+    POST /evaluate  {"state": ..., "questions": {id: {type, instructions, criteria?}}}
+                 -> {"answers": {...}, "usage": {...}, "warnings": [...], "providerMetadata": {"laya": {...}}}
+    GET  /health -> {"status": "ok" | "warming", ...}
+
+Question types use the gateway names: `boolean` (Laya calls it `noul`), `choice`, `score`.
+
+The server binds to loopback only, never logs state or question text, and keeps
+the model behind one lock because a Core ML model is not safe to call concurrently.
+"""
+
+import argparse
+import json
+import os
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 47823
+MAX_BODY_BYTES = 1024 * 1024
+# Measured on an M3 Max with 10-question preflights: no floor p50 411 ms (one shape switch per request),
+# 256 -> 124 ms, 384 -> 155 ms. Requests longer than the floor still work; they just pay the switch.
+DEFAULT_MIN_LENGTH = 256
+DEFAULT_MODEL_DIR = Path(__file__).resolve().parent / "models" / "laya-multilingual-coreml"
+
+
+class BadRequest(ValueError):
+    pass
+
+
+def setting(name, default):
+    """`MU_<name>`, or the `KYRN_<name>` spelling from before the rename."""
+    return os.environ.get(f"MU_{name}") or os.environ.get(f"KYRN_{name}") or default
+
+
+def _text(value):
+    """Laya serializes non-string instructions with ASCII escapes, which mangles CJK. Do it here instead."""
+    return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
+
+
+def to_laya_question(question_id, question):
+    if not isinstance(question, dict):
+        raise BadRequest(f'Question "{question_id}" must be an object')
+    kind = question.get("type")
+    if "instructions" not in question:
+        raise BadRequest(f'Question "{question_id}" is missing instructions')
+    instructions = _text(question["instructions"])
+    criteria = question.get("criteria")
+
+    if kind == "boolean":
+        out = {"type": "noul", "instructions": instructions}
+        if isinstance(criteria, dict):
+            kept = {key: criteria[key] for key in ("false", "true") if criteria.get(key) not in (None, "")}
+            if kept:
+                out["criteria"] = kept
+        return out
+    if kind == "choice":
+        if not isinstance(criteria, dict) or not criteria:
+            raise BadRequest(f'Choice question "{question_id}" needs a non-empty criteria object')
+        return {"type": "choice", "instructions": instructions, "criteria": dict(criteria)}
+    if kind == "score":
+        if not isinstance(criteria, list) or len(criteria) < 2:
+            raise BadRequest(f'Score question "{question_id}" needs at least two levels')
+        return {
+            "type": "score",
+            "instructions": instructions,
+            "criteria": ["" if level is None else level for level in criteria],
+        }
+    raise BadRequest(f'Question "{question_id}" has unknown type {kind!r}; expected boolean, choice or score')
+
+
+def from_laya_answer(answer):
+    kind = answer["type"]
+    if kind == "noul":
+        return {"type": "boolean", "probability": answer["noul"]}
+    if kind == "choice":
+        return {"type": "choice", "choice": answer["choice"], "probabilities": answer["probabilities"]}
+    return {"type": "score", "score": answer["score"], "probabilities": answer["probabilities"]}
+
+
+class LocalJudge:
+    def __init__(self, model_dir, compute_units=None, min_length=0):
+        import laya_coreml as laya
+
+        self.model_dir = Path(model_dir)
+        self.min_length = min_length
+        started = time.perf_counter()
+        self.agent = laya.load(str(self.model_dir), local_files_only=True, compute_units=compute_units)
+        self.load_seconds = time.perf_counter() - started
+        self.lock = threading.Lock()
+        self.warm = False
+        self.requests = 0
+        self.started_at = time.time()
+        manifest = getattr(self.agent, "manifest", {})
+        self.model_id = "laya:" + str(manifest.get("repository") or self.model_dir.name)
+        self.shape = getattr(self.agent, "shape", {})
+
+    def uniform_shape(self, state, laya_questions):
+        """Pad every question of a request to one input length.
+
+        Core ML re-specializes the GPU program whenever the enumerated length
+        changes (about 350 ms on an M3 Max, against 40 ms for a steady shape),
+        and the questions of one request differ in length because their
+        options do. One bucket per request removes those switches; the floor
+        keeps consecutive requests on the same bucket too.
+        """
+        lengths = self.agent.shape.get("lengths")
+        if not lengths:
+            return self.agent.shape
+        items, _internal = self.agent.prepare(state, laya_questions)
+        longest = max(max(len(item["ids"]) for item in items), self.min_length)
+        target = next((n for n in lengths if n >= longest), lengths[-1])
+        return dict(self.agent.shape, lengths=[target])
+
+    def diagnostics(self, state, laya_questions):
+        """Report what the token budget silently cut. Never raises."""
+        warnings = []
+        try:
+            from laya_coreml.common import build_prefix, render_options, serialize_state
+
+            tok, cfg = self.agent.tok, self.agent.cfg
+            max_len, head_max_len = cfg.get("max_len", 512), cfg.get("head_max_len", 192)
+            state_tokens = len(tok(serialize_state(state).replace(tok.mask_token, " "))["input_ids"])
+            for question_id, definition in laya_questions.items():
+                q = self.agent._to_internal(definition)
+                prefix, _markers = build_prefix(tok, q, head_max_len)
+                full_head = len(tok("%s question: %s" % (q["t"], q["ins"]))["input_ids"])
+                full_options = sum(1 + len(tok(" " + option)["input_ids"]) for option in render_options(q))
+                full_prefix = 1 + full_head + 1 + full_options + 1
+                if len(prefix) < full_prefix:
+                    warnings.append(
+                        {
+                            "type": "question_truncated",
+                            "questionId": question_id,
+                            "message": f"question and options need {full_prefix} tokens, {len(prefix)} were kept",
+                        }
+                    )
+                room = max(0, max_len - len(prefix) - 1)
+                if state_tokens > room:
+                    warnings.append(
+                        {
+                            "type": "state_truncated",
+                            "questionId": question_id,
+                            "message": f"state has {state_tokens} tokens, only the first {room} were read",
+                        }
+                    )
+        except Exception as error:  # noqa: BLE001 - diagnostics are best effort
+            warnings.append({"type": "diagnostics_failed", "message": type(error).__name__})
+        return warnings
+
+    def evaluate(self, state, questions):
+        if not isinstance(questions, dict) or not questions:
+            raise BadRequest("questions must be a non-empty object keyed by question id")
+        if not isinstance(state, (str, dict, list)):
+            raise BadRequest("state must be a string, an object or an array")
+        laya_questions = {qid: to_laya_question(qid, q) for qid, q in questions.items()}
+
+        started = time.perf_counter()
+        with self.lock:
+            waited = time.perf_counter() - started
+            original_shape = self.agent.shape
+            try:
+                self.agent.shape = self.uniform_shape(state, laya_questions)
+                result = self.agent.predict(state, laya_questions)
+            except ValueError as error:
+                raise BadRequest(str(error)) from error
+            finally:
+                self.agent.shape = original_shape
+            self.requests += 1
+        elapsed = time.perf_counter() - started
+
+        metadata = {
+            qid: {"confidence": answer.get("confidence"), "actProbability": answer.get("action", {}).get("act_probability")}
+            for qid, answer in result["answers"].items()
+        }
+        usage = result.get("usage", {})
+        print(
+            f"[{time.strftime('%H:%M:%S')}] evaluate questions={len(questions)} "
+            f"tokens={usage.get('input_tokens', 0)} ms={elapsed * 1000:.0f} waited_ms={waited * 1000:.0f}",
+            file=sys.stderr,
+            flush=True,
+        )
+        return {
+            "answers": {qid: from_laya_answer(answer) for qid, answer in result["answers"].items()},
+            "usage": {"inputTokens": usage.get("input_tokens", 0), "outputTokens": 0},
+            "warnings": self.diagnostics(state, laya_questions),
+            "providerMetadata": {"laya": {"model": self.model_id, "latencyMs": round(elapsed * 1000, 1), "answers": metadata}},
+        }
+
+    def warm_up(self):
+        """Core ML compiles each enumerated input length on first use (seconds). Pay that before real traffic."""
+        lengths = list(self.shape.get("lengths") or [self.shape.get("default_length", 128)])
+        # Finish on the bucket real traffic starts in, so the first request does not pay a shape switch.
+        resting = next((n for n in lengths if n >= self.min_length), lengths[-1])
+        lengths = [n for n in lengths if n != resting] + [resting]
+        question = {"warm": {"type": "noul", "instructions": "Is this a warm-up request?"}}
+        started = time.perf_counter()
+        for length in lengths:
+            # "a " is roughly one token; the prefix takes about 30, so this lands in the bucket for `length`.
+            state = "a " * max(1, length - 40)
+            try:
+                with self.lock:
+                    original_shape = self.agent.shape
+                    self.agent.shape = dict(original_shape, lengths=[length])
+                    try:
+                        self.agent.predict(state, question)
+                    finally:
+                        self.agent.shape = original_shape
+            except Exception as error:  # noqa: BLE001 - a failed warm-up only costs latency later
+                print(f"warm-up for length {length} failed: {type(error).__name__}", file=sys.stderr, flush=True)
+        self.warm = True
+        print(f"warm-up finished in {time.perf_counter() - started:.1f}s ({len(lengths)} shapes)", file=sys.stderr, flush=True)
+
+    def health(self):
+        return {
+            "status": "ok" if self.warm else "warming",
+            "model": self.model_id,
+            "computeUnits": getattr(self.agent, "compute_units", None),
+            "maxLength": self.shape.get("max_length"),
+            "maxOptions": self.shape.get("max_options"),
+            "batchSize": self.shape.get("batch_size"),
+            "requests": self.requests,
+            "uptimeSeconds": round(time.time() - self.started_at),
+            "pid": os.getpid(),
+        }
+
+
+def make_handler(judge):
+    class Handler(BaseHTTPRequestHandler):
+        server_version = "mu-local-judge/0.1"
+
+        def log_message(self, *_args):
+            pass
+
+        def send_json(self, status, payload):
+            body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
+            self.send_response(status)
+            self.send_header("Content-Type", "application/json; charset=utf-8")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+
+        def send_failure(self, status, message):
+            self.send_json(status, {"error": {"message": message}})
+
+        def do_GET(self):
+            if self.path == "/health":
+                self.send_json(200, judge.health())
+            else:
+                self.send_failure(404, "Not found")
+
+        def do_POST(self):
+            if self.path != "/evaluate":
+                return self.send_failure(404, "Not found")
+            # Browsers always send Origin on cross-site POSTs; nothing legitimate calls this from a web page.
+            if self.headers.get("Origin"):
+                return self.send_failure(403, "Browser requests are not allowed")
+            if "application/json" not in (self.headers.get("Content-Type") or ""):
+                return self.send_failure(415, "Content-Type must be application/json")
+            try:
+                length = int(self.headers.get("Content-Length") or 0)
+            except ValueError:
+                return self.send_failure(400, "Invalid Content-Length")
+            if length <= 0 or length > MAX_BODY_BYTES:
+                return self.send_failure(413, f"Body must be between 1 and {MAX_BODY_BYTES} bytes")
+            try:
+                payload = json.loads(self.rfile.read(length))
+            except (ValueError, UnicodeDecodeError):
+                return self.send_failure(400, "Body is not valid JSON")
+            if not isinstance(payload, dict):
+                return self.send_failure(400, "Body must be a JSON object")
+            try:
+                self.send_json(200, judge.evaluate(payload.get("state"), payload.get("questions")))
+            except BadRequest as error:
+                self.send_failure(400, str(error)[:300])
+            except Exception as error:  # noqa: BLE001 - the kernel fails open on a 500
+                print(f"evaluate failed: {type(error).__name__}: {error}", file=sys.stderr, flush=True)
+                self.send_failure(500, f"Local judge failed: {type(error).__name__}")
+
+    return Handler
+
+
+def main():
+    parser = argparse.ArgumentParser(description="mu local judge sidecar (Laya on Core ML)")
+    parser.add_argument("--model-dir", default=setting("LOCAL_JUDGE_MODEL", str(DEFAULT_MODEL_DIR)))
+    parser.add_argument("--host", default=DEFAULT_HOST)
+    parser.add_argument("--port", type=int, default=int(setting("LOCAL_JUDGE_PORT", DEFAULT_PORT)))
+    parser.add_argument("--compute-units", choices=["all", "cpu", "cpu_gpu", "cpu_ne"])
+    parser.add_argument("--no-warmup", action="store_true")
+    parser.add_argument(
+        "--min-length",
+        type=int,
+        default=int(setting("LOCAL_JUDGE_MIN_LENGTH", DEFAULT_MIN_LENGTH)),
+        help="Smallest input length to pad to. A floor above the usual request size avoids shape switches between requests.",
+    )
+    args = parser.parse_args()
+
+    if args.host not in ("127.0.0.1", "localhost", "::1"):
+        parser.error("The local judge has no authentication and only binds to loopback")
+
+    judge = LocalJudge(args.model_dir, compute_units=args.compute_units, min_length=args.min_length)
+    server = ThreadingHTTPServer((args.host, args.port), make_handler(judge))
+    server.daemon_threads = True
+    print(
+        f"mu local judge: {judge.model_id} loaded in {judge.load_seconds:.1f}s, "
+        f"listening on http://{args.host}:{args.port}",
+        file=sys.stderr,
+        flush=True,
+    )
+    if args.no_warmup:
+        judge.warm = True
+    else:
+        threading.Thread(target=judge.warm_up, daemon=True).start()
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        server.server_close()
+
+
+if __name__ == "__main__":
+    main()

package/package.json

@@ -0,0 +1,45 @@
+{
+	"name": "mu-agent",
+	"version": "0.1.0",
+	"description": "mu: a coding agent that thinks before it acts. A small, fast judge makes the routine calls; the big model keeps its attention for the work. Built on pi.",
+	"type": "module",
+	"bin": {
+		"mu": "kyrn/bin/mu.mjs"
+	},
+	"piConfig": {
+		"name": "mu",
+		"configDir": ".mu"
+	},
+	"keywords": [
+		"coding-agent",
+		"ai",
+		"llm",
+		"cli",
+		"terminal",
+		"agent",
+		"judge",
+		"pi"
+	],
+	"homepage": "https://github.com/qybaihe/MU",
+	"repository": {
+		"type": "git",
+		"url": "git+https://github.com/qybaihe/MU.git"
+	},
+	"bugs": {
+		"url": "https://github.com/qybaihe/MU/issues"
+	},
+	"license": "MIT",
+	"engines": {
+		"node": ">=22.19.0"
+	},
+	"dependencies": {
+		"@earendil-works/chord": "0.86.0",
+		"@silvia-odwyer/photon-node": "0.3.4",
+		"jiti": "2.7.0"
+	},
+	"muBuild": {
+		"pi": "0.86.0",
+		"judge": "0.1.0",
+		"commit": "d5e481d01"
+	}
+}