Run Muse Spark Open locally with llama.cpp: hardware, quants and setup

TL;DR: Muse Spark Open is Meta's 816B-parameter open-weight mixture-of-experts model, with about 42B parameters active per token. You can run it on one machine if RAM plus VRAM, or Apple unified memory, adds up to roughly 230 GiB or more. With 256 GiB, use Unsloth's UD-IQ2_XXS GGUF (225.6 GB) in llama.cpp. With 512 GiB, use UD-Q4_K_XL (505.5 GB). Run the included planner first: it reads your memory and context length and prints the exact download and llama-server commands. Then send a streamed request with reasoning_effort set to low. This guide is source-reviewed. Both scripts were executed on a CPU; we have not yet run the model itself.

1. Know what you are about to run

Meta Superintelligence Labs released Muse Spark Open on 15 October 2026. It reads text, images, video and audio, writes text, and has a context window of 1,048,576 tokens. Meta trained it for agentic coding: reading a repository, calling tools and correcting its own mistakes. The facts below come from the model card and the Hugging Face launch post.

A mixture-of-experts (MoE) layer holds many small feed-forward networks, the experts. A router picks a few of them for each token. You must hold every expert in memory, but each token only runs through the ones the router picks.

The important consequence: memory follows the 816B, and speed follows the 42B. A model with 42B active parameters runs at roughly the speed of a 42B dense model, but you have to store all 816B. That is why a model this large runs on a Mac or a big-memory workstation at all, and why it needs so much memory.

The official checkpoints are BF16 and NVFP4. Meta's deployment guide serves NVFP4 on four B300 GPUs and BF16 on eight. To run the model locally you need GGUF files: the format llama.cpp loads, with weights compressed to between about 2 and 8 bits. Two teams publish them. Unsloth ships eight "dynamic" quantizations. ggml-org, the llama.cpp maintainers, ships Q2_K, Q4_K_S, Q8_0 and BF16. When we checked on 19 October 2026, neither GGUF repository required access approval. The official BF16 repository did.

The license still applies to GGUF copies. Its notable terms, as summarized on the model card: show "Built with Muse Spark Open" in products that use it; a model-as-a-service offering with more than $20M revenue in 12 months needs a separate agreement with Meta; and no users under 18. Read the full license file in the repository you download.

Check: you can explain why an 816B model can produce tokens at the speed of a 42B model while still needing 816B worth of memory.

2. Decide whether your machine can run it

Add three things and compare the total with your memory:

1. The GGUF file size. Fixed by the quant you pick.
2. The KV cache. The model stores keys and values for every token in the context.
3. A reserve. The operating system, the runtime and its compute buffers. The planner assumes 12 GiB. Adjust it for your machine.

The KV-cache formula, using the published architecture:

values per token = 2 (K and V) × 62 layers × 16 KV heads × 64 dims = 126,976
bytes per token  = 126,976 × 2 bytes (f16)                         = 253,952 B ≈ 248 KiB
32,768 tokens    = 32,768 × 253,952 B = 8.32 GB                    = 7.75 GiB

Hugging Face's launch post gives the same per-token figure and about 266 GB for the full 1M window. That agrees with this formula: 1,048,576 × 253,952 B = 266.3 GB.

q8_0 stores 32 values in 34 bytes, so it uses 53% of the f16 size. It is a llama.cpp KV-cache option (--cache-type-k, --cache-type-v). It can change output slightly, so compare answers on your own prompts before relying on it.

Now the weights. These are the byte totals Hugging Face reported for each quant on 19 October 2026, with the full total at a 32K f16 context:

Downloads are in GB (10⁹ bytes) because that is how Hugging Face lists them. Memory is in GiB (2³⁰ bytes) because a "256 GB" Mac or a 256 GB DIMM kit actually holds 256 GiB, which is 274.9 GB. That unit gap is why a 225.6 GB download fits comfortably on a 256 GB machine.

Pick your path:

Run the planner. It needs only Python 3.10 or newer and downloads nothing:

mkdir muse-local && cd muse-local
# save plan_memory.py and smoke.py from Project files here
python3 plan_memory.py --ram 256
python3 plan_memory.py --ram 128 --vram 48 --ctx 65536 --kv q8_0

--ram is system RAM, or unified memory on a Mac. --vram is the total across discrete GPUs. Add --use agent to skip quants below UD-IQ2_S, and --vision if you plan to send images or audio.

Check: the planner prints a table and a line starting Largest quant that fits:, followed by a hf download command and a llama-server command. If it prints Nothing fits, lower --ctx, try --kv q8_0, or go to step 9.

3. Install llama.cpp

llama.cpp runs the GGUF on CPU, NVIDIA (CUDA), AMD (ROCm), Vulkan and Apple Metal. Support for Muse Spark Open was merged on 15 October 2026, so use a build from that day or later. Meta's guide installs the prebuilt llama tool:

curl -LsSf https://llama.app/install.sh | sh

To build from source with CUDA instead, follow Unsloth's Muse Spark Open guide. In short:

git clone https://github.com/ggml-org/llama.cpp
cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=ON
cmake --build llama.cpp/build --config Release -j --target llama-server llama-cli

Use -DGGML_CUDA=OFF for a CPU-only or Apple build. Metal is enabled by default on macOS. You also need the Hugging Face CLI for resumable downloads:

python3 -m pip install -U huggingface_hub

Check: llama-server --version (or ./llama.cpp/build/bin/llama-server --version) prints a build from 15 October 2026 or later, and hf --help works.

Prefer a desktop app? Unsloth Desktop wraps the same llama.cpp engine. Install it, search its Model Hub for "Muse Spark Open", download a quant the planner approved, and skip to step 6. Its API server listens on the port Unsloth shows you; pass that URL to smoke.py --url.

4. Download the weights

Check free disk space first. You need the download size plus some slack: about 240 GB for UD-IQ2_XXS. Then run the hf download line the planner printed, for example:

df -h .
hf download unsloth/Muse-Spark-Open-816B-A42B-GGUF \
  --include "UD-IQ2_XXS/*" --local-dir muse-gguf

--include limits the download to one quant folder. Without it you would fetch every quant: more than 2.5 TB. If the connection drops, run the same command again. Files that finished downloading are skipped.

These times are the file size divided by the link rate, so they are a best case. Real downloads are usually slower.

Check: ls muse-gguf/UD-IQ2_XXS lists six files ending -00001-of-00006.gguf through -00006-of-00006.gguf. llama.cpp opens the first shard and finds the others itself; do not merge them.

5. Start the server

Run the llama-server command from the planner. For the 256 GiB case:

llama-server \
  --model muse-gguf/UD-IQ2_XXS/Muse-Spark-Open-816B-A42B-UD-IQ2_XXS-00001-of-00006.gguf \
  --alias muse-spark-open \
  --ctx-size 32768 \
  --host 127.0.0.1 --port 8080 \
  --temp 1.0 --top-p 1.0 --top-k 0 --min-p 0.0 \
  --chat-template-kwargs '{"reasoning_effort":"low"}'

You do not need to say which layers go to the GPU. Current llama-server fits weights into GPU memory automatically (--fit on is the default) and leaves the rest in system RAM. For manual control, --cpu-moe keeps all expert weights in RAM, and --n-cpu-moe N keeps the experts of the first N layers there. Experts hold about 97% of the parameters, so a 24 to 48 GB GPU mostly speeds up attention and prompt processing. The flag reference is the llama-server README.

Loading takes a while: llama.cpp maps 225.6 GB of files into memory. From an NVMe drive that reads 3.5 GB/s, that is at least 64 seconds before any warm-up.

Check: from a second terminal, wait until the health endpoint returns {"status":"ok"}. It returns HTTP 503 with Loading model until then:

curl -s http://127.0.0.1:8080/health

6. Send a test request

First a plain request, so you can see the raw response shape:

curl -s http://127.0.0.1:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"muse-spark-open","messages":[{"role":"user","content":"What is 17 * 23?"}]}'

The response holds the thinking in reasoning_content and the answer in content. Muse Spark Open always reasons. Meta's guide says the template rejects "reasoning_effort": "none", so low is the cheapest setting.

Then run the smoke test. It streams the request, so a long reasoning trace cannot hit a fixed timeout, and it checks that the stream ended normally:

python3 smoke.py
python3 smoke.py --effort medium

Check: smoke.py prints JSON with "ok": true, "finish_reason": "stop", a non-empty answer and "answer_correct": true (the answer contains 391). reasoning_chars shows how much the model thought before answering. If you see hit max_tokens, the thinking used the whole budget: raise --max-tokens or lower --effort.

Muse Spark Open stops on <|eot|>, its end-of-turn token. <|eom|> only ends one message inside a turn. llama-server reads this from the GGUF; it matters only if you write your own stop rules.

7. Add images and audio (optional)

The language-model GGUF reads text only. For images and audio, download the projector and pass it with --mmproj:

hf download unsloth/Muse-Spark-Open-816B-A42B-GGUF \
  --include "mmproj-Q8_0.gguf" --local-dir muse-gguf
# add to the llama-server command:
#   --mmproj muse-gguf/mmproj-Q8_0.gguf

The Q8_0 projector is 2.2 GB. Run the planner with --vision to include it. Send an image as an OpenAI-style image_url content part (a URL or a data:image/png;base64,... string). Audio uses an input_audio part. Audio becomes 12.5 tokens per second of sound: 125 tokens for a 10-second clip. The model card says video with a short side under about 448 pixels may fail.

Check: a request with one image and the text "What is in this image?" returns a description of that image.

8. Engineer's notes: what to measure and how to compare

This section is our own analysis of the published numbers. Every figure is either computed from sources linked here or quoted from a named measurement. We have not run the model ourselves.

The real bits per weight

Quant names describe the most common format inside the file, not the average. Dividing each file's bytes by 816B parameters gives the actual average:

Two things to take from this. First, the BF16 GGUF works out to 16.08 bits per weight. That confirms the 816B parameter count from the files alone (the small excess is metadata and a few higher-precision tensors). Second, "2-bit" files differ by a third in size: 2.21 against 2.98 bits. Dynamic quants keep sensitive tensors at higher precision. When you compare two quants, compare file size, not the name.

The last column multiplies the average bits by Unsloth's figure of about 42.7B active parameters. That assumes the active tensors are compressed like the average. Unsloth's dynamic quants keep some always-active tensors at higher precision, so treat the column as a lower bound.

A speed ceiling you can compute

Generating a token is limited mainly by memory bandwidth: every active weight is read once per token. So:

tokens/s ceiling ≈ memory bandwidth ÷ bytes read per token

Check it against the two measurements Hugging Face published on a Mac Studio with M3 Ultra and 512 GB. Apple lists that machine at over 800 GB/s:

Reaching 35 to 50% of the bandwidth ceiling is normal for MoE decoding at batch size 1. Routing, small expert matrix multiplies and synchronization all cost time. Use the ceiling to rule hardware out, not to predict an exact speed. Do not conclude from this table that MLX is faster than llama.cpp. The two runs used different quants, and Hugging Face did not publish the llama.cpp prompt and output lengths.

Apply the same arithmetic to your own machine before you download 225 GB:

DDR5 bandwidth here is channels × 8 bytes × transfer rate: 2 × 8 × 5.6 GT/s = 89.6 GB/s. With --cpu-moe on a desktop with one GPU, nearly every byte read per token still comes from system RAM, so the desktop row applies even with a fast GPU. Expect a fraction of that ceiling in practice.

At long context, the cache dominates

Every new token also reads the whole KV cache. All 62 layers use full attention, so that read grows linearly with context. For UD-IQ2_XXS at 800 GB/s:

At 128K tokens of context, the cache read is nearly three times the weight read. This is the most useful number in this section: a long agentic session slows down as it goes, even though nothing about the model changed. Quantizing the cache to q8_0 roughly halves that column. Measure the effect on your own sessions rather than assuming it.

How Muse Spark Open's cache compares

The same per-token formula, applied to the published configuration of other large open MoE models:

DeepSeek-V3's multi-head latent attention (MLA) caches one compressed vector per layer, about 3.6 times less than Muse Spark Open per token. gpt-oss limits half its layers to a 128-token window. Muse Spark Open uses neither technique. Its 16 KV heads keep the cache moderate for its size, but long-context memory and speed cost more than on DeepSeek-V3 or Kimi K2, which uses the same MLA layout as DeepSeek-V3. Configurations: Qwen3-235B-A22B, DeepSeek-V3, gpt-oss-120b. For the method, see KV cache sizing.

Where the official instructions disagree

We found these inconsistencies while reviewing the launch documents. Each one can cost you an hour:

What to watch on your own run

Record these four things on your first run, and keep them with the quant name and context size:

1. Swap. Run free -h (Linux) or check Activity Monitor's memory pressure (macOS) during generation. If the system swaps, speed falls by orders of magnitude. Use a smaller quant or context; do not wait it out.
2. Generation speed against the ceiling. llama-server logs prompt and generation speed after each request. Below 20% of your ceiling, look for swap, CPU threads competing with other work, or experts split across slow links.
3. Reasoning share. Run smoke.py at each effort level on one of your real prompts and compare reasoning_chars and seconds_total. At xhigh, a short answer can come after a very long trace.
4. Speed at long context. Repeat a request after a long conversation. The table above predicts the slowdown; measuring it tells you how long a local agent session stays usable.

9. Other ways to run it

Exact server commands for vLLM, SGLang, TensorRT-LLM and MLX are in Meta's deployment guide. The GPU routes serve an OpenAI-compatible API, so smoke.py --url http://HOST:8000 --model meta-models/Muse-Spark-Open-816B-A42B-NVFP4 runs the same check against them. Hosted APIs need an API key; set MUSE_API_KEY and pass the provider's base URL and model string. For sizing a multi-GPU deployment, see designing inference for a 1T-parameter MoE and expert parallelism.

10. Troubleshoot the first failure

11. Clean up

Stop the server with Ctrl+C in its terminal. Nothing else keeps running. The weights stay in muse-gguf/ so you can restart without downloading again. To reclaim the space, check what you are deleting first:

du -sh muse-gguf
rm -r muse-gguf/UD-IQ2_XXS

hf download --local-dir may also have cached metadata under muse-gguf/.cache. Remove the whole muse-gguf folder when you are finished with the model.

Frequently asked questions

Can I run Muse Spark Open on 128 GB of RAM? Not well. The smallest GGUF, UD-IQ1_S, is 175.6 GB. llama.cpp can memory-map the rest from disk, but every token then reads gigabytes from the drive, and generation becomes extremely slow. Use a hosted provider instead.

How much VRAM do I need? None is required. llama.cpp can run entirely from system RAM. VRAM counts toward the total, and a GPU speeds up prompt processing, but the experts usually stay in RAM on a single consumer GPU.

Which quant should I pick? The largest one the planner says fits, at the context length you need. For coding agents and tool calls, Unsloth recommends UD-IQ2_S (275.5 GB) or larger.

Can I turn reasoning off? No. Use reasoning_effort: "low" for the shortest thinking.

Can I use the full 1M-token context locally? Only with a lot of spare memory. The f16 cache for 1M tokens is 248 GiB on top of the weights, or about 132 GiB with q8_0.

Is Muse Spark Open free for commercial use? The Muse Model License allows it with conditions, including the "Built with Muse Spark Open" notice and a separate agreement for model-as-a-service offerings above $20M revenue in 12 months. Read the license itself; this summary is not legal advice.

How does it compare with Muse Glimmer? Muse Glimmer 30B is the small dense model of the same family, with the same tokenizer. It fits a single workstation GPU. Muse Spark Open is the full-size MoE.

Check your understanding

Your machine has 256 GiB of RAM. Why does UD-IQ2_XXS fit at a 32K context, and UD-IQ2_S does not? UD-IQ2_XXS is 210.1 GiB of weights, plus 7.75 GiB of cache and the 12 GiB reserve: 229.8 GiB. UD-IQ2_S is 256.6 GiB of weights alone, already over the limit before the cache.

A long coding session gets slower as it goes, with no change to the model. What changed? The context. Each new token reads the KV cache for every earlier token. At 128K tokens that read is larger than the weight read, so the speed ceiling falls to under a third of its short-prompt value.

Complete project source

These are the same files offered under Project files. The planner and smoke client have CPU tests in this repository; the model run remains unmeasured.

plan_memory.py

"""Pick a Muse Spark Open GGUF that fits your memory, then print the commands.

Runs anywhere with Python 3.10+. It downloads nothing and loads no model.
File sizes are the byte totals Hugging Face reported for each quant folder
on 2026-10-19. The KV-cache size follows from the published architecture.
Everything else (the reserve for the OS and runtime) is an assumption you
can change with --reserve.

  python3 plan_memory.py --ram 256
  python3 plan_memory.py --ram 128 --vram 24 --ctx 65536 --kv q8_0
  python3 plan_memory.py --ram 512 --use agent --json
"""
import argparse
import json
import sys

GIB = 2**30
UNSLOTH = "unsloth/Muse-Spark-Open-816B-A42B-GGUF"
GGML = "ggml-org/Muse-Spark-Open-816B-A42B-GGUF"

# (quant, repository, bytes, shards). Smallest first.
QUANTS = [
    ("UD-IQ1_S", UNSLOTH, 175_639_574_392, 5),
    ("UD-IQ1_M", UNSLOTH, 208_925_571_064, 6),
    ("UD-IQ2_XXS", UNSLOTH, 225_568_569_304, 6),
    ("UD-IQ2_S", UNSLOTH, 275_497_564_280, 7),
    ("Q2_K", GGML, 304_080_474_208, 22),
    ("UD-Q2_K_XL", UNSLOTH, 311_225_649_368, 8),
    ("UD-IQ3_XXS", UNSLOTH, 340_350_896_376, 8),
    ("UD-Q3_K_XL", UNSLOTH, 401_483_491_832, 10),
    ("Q4_K_S", GGML, 470_069_227_350, 22),
    ("UD-Q4_K_XL", UNSLOTH, 505_502_231_288, 12),
    ("Q8_0", GGML, 871_648_188_873, 10),
]
# Unsloth's guide: UD-IQ2_S is the smallest quant to use for agentic work.
AGENT_MINIMUM = "UD-IQ2_S"
VISION_PROJECTOR = 2_183_868_768  # mmproj-Q8_0.gguf, needed for image/audio input

# 62 layers x 16 KV heads x 64 dims, for K and V: 126,976 values per token.
LAYERS, KV_HEADS, HEAD_DIM = 62, 16, 64
VALUES_PER_TOKEN = 2 * LAYERS * KV_HEADS * HEAD_DIM
# llama.cpp block formats: q8_0 stores 32 values in 34 bytes, q4_0 in 18.
BYTES_PER_VALUE = {"f16": 2.0, "bf16": 2.0, "q8_0": 34 / 32, "q4_0": 18 / 32}
MAX_CONTEXT = 1_048_576


def kv_bytes(ctx: int, kv: str) -> int:
    return round(ctx * VALUES_PER_TOKEN * BYTES_PER_VALUE[kv])


def shard_path(quant: str, repo: str, shards: int) -> str:
    name = f"Muse-Spark-Open-816B-A42B-{quant}-00001-of-{shards:05d}.gguf"
    return f"{quant}/{name}" if repo == UNSLOTH else name


def plan(memory_gib: float, ctx: int, kv: str, reserve_gib: float, vision: bool, use: str):
    if not 1 <= ctx <= MAX_CONTEXT:
        raise ValueError(f"--ctx must be between 1 and {MAX_CONTEXT}")
    if memory_gib <= 0 or reserve_gib < 0:
        raise ValueError("memory must be positive and the reserve cannot be negative")
    cache = kv_bytes(ctx, kv)
    extra = cache + reserve_gib * GIB + (VISION_PROJECTOR if vision else 0)
    names = [q[0] for q in QUANTS]
    rows = []
    for quant, repo, size, shards in QUANTS:
        need = size + extra
        rows.append({
            "quant": quant, "repo": repo, "download_gb": round(size / 1e9, 1),
            "need_gib": round(need / GIB, 1), "fits": need <= memory_gib * GIB,
            "agent_ok": names.index(quant) >= names.index(AGENT_MINIMUM),
            "first_shard": shard_path(quant, repo, shards),
        })
    usable = [r for r in rows if r["fits"] and (use == "chat" or r["agent_ok"])]
    return {
        "memory_gib": memory_gib, "ctx": ctx, "kv": kv, "kv_cache_gib": round(cache / GIB, 2),
        "reserve_gib": reserve_gib, "vision": vision, "use": use, "rows": rows,
        "choice": usable[-1] if usable else None,
    }


def commands(result) -> list[str]:
    c = result["choice"]
    if c is None:
        return []
    include = f"{c['quant']}/*" if c["repo"] == UNSLOTH else f"*-{c['quant']}-*"
    model = f"muse-gguf/{c['first_shard']}"
    serve = [
        "llama-server", f"--model {model}", "--alias muse-spark-open",
        f"--ctx-size {result['ctx']}", "--host 127.0.0.1 --port 8080",
        "--temp 1.0 --top-p 1.0 --top-k 0 --min-p 0.0",
        "--chat-template-kwargs '{\"reasoning_effort\":\"low\"}'",
    ]
    if result["kv"] != "f16":
        serve.insert(4, f"--cache-type-k {result['kv']} --cache-type-v {result['kv']}")
    lines = [f'hf download {c["repo"]} --include "{include}" --local-dir muse-gguf']
    if result["vision"]:
        projector = "mmproj-Q8_0.gguf" if c["repo"] == UNSLOTH else "mmproj-Muse-Spark-Open-816B-A42B-Q8_0.gguf"
        lines.append(f'hf download {c["repo"]} --include "{projector}" --local-dir muse-gguf')
        serve.insert(2, f"--mmproj muse-gguf/{projector}")
    lines.append(" \\\n  ".join(serve))
    return lines


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--ram", type=float, required=True, help="system RAM or unified memory, GiB")
    p.add_argument("--vram", type=float, default=0, help="total GPU memory on discrete GPUs, GiB")
    p.add_argument("--ctx", type=int, default=32768, help="context tokens to reserve KV cache for")
    p.add_argument("--kv", choices=sorted(BYTES_PER_VALUE), default="f16", help="KV-cache type")
    p.add_argument("--reserve", type=float, default=12, help="GiB kept for OS, runtime and buffers")
    p.add_argument("--vision", action="store_true", help="add the image/audio projector")
    p.add_argument("--use", choices=["chat", "agent"], default="chat")
    p.add_argument("--json", action="store_true")
    a = p.parse_args(argv)
    try:
        result = plan(a.ram + a.vram, a.ctx, a.kv, a.reserve, a.vision, a.use)
    except ValueError as error:
        p.error(str(error))
    result["commands"] = commands(result)
    if a.json:
        print(json.dumps(result, indent=2))
        return 0 if result["choice"] else 1
    print(f"Memory {result['memory_gib']:g} GiB | KV cache {result['kv_cache_gib']} GiB "
          f"for {a.ctx:,} tokens ({a.kv}) | reserve {a.reserve:g} GiB")
    print(f"{'quant':<12}{'download':>11}{'needs':>11}  fits")
    for r in result["rows"]:
        mark = "yes" if r["fits"] else "no"
        if r["fits"] and a.use == "agent" and not r["agent_ok"]:
            mark = "chat only"
        print(f"{r['quant']:<12}{r['download_gb']:>8} GB{r['need_gib']:>7} GiB  {mark}")
    if not result["choice"]:
        print("\nNothing fits. Lower --ctx, use --kv q8_0, add memory, or use a hosted provider.")
        return 1
    print(f"\nLargest quant that fits: {result['choice']['quant']}\n")
    print("\n".join(result["commands"]))
    return 0


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

smoke.py

"""Check that a local Muse Spark Open server answers a real chat request.

Works with llama-server, Unsloth and any OpenAI-compatible endpoint. Uses only
the Python standard library. The request streams, because Muse Spark Open
always reasons before it answers and a long trace can outlast a fixed timeout.

  python3 smoke.py
  python3 smoke.py --url http://127.0.0.1:8888 --effort medium
  MUSE_API_KEY=... python3 smoke.py --model muse-spark-open
"""
import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request

EFFORTS = ("low", "medium", "high", "xhigh")
PROMPT = "What is 17 * 23? Reply with the number only."


def request(url, key, body=None, timeout=30):
    headers = {"Content-Type": "application/json"}
    if key:
        headers["Authorization"] = f"Bearer {key}"
    data = None if body is None else json.dumps(body).encode()
    return urllib.request.urlopen(urllib.request.Request(url, data, headers), timeout=timeout)


def events(response):
    """Yield each JSON chunk of a server-sent event stream."""
    for raw in response:
        line = raw.decode("utf-8").strip()
        if not line.startswith("data:"):
            continue
        payload = line[5:].strip()
        if payload == "[DONE]":
            return
        yield json.loads(payload)


def run(url, model, effort, max_tokens, key, timeout):
    base = url.rstrip("/")
    result = {"url": base, "model": model, "reasoning_effort": effort, "ok": False}
    try:
        with request(f"{base}/v1/models", key, timeout=10) as r:
            result["served_models"] = [m["id"] for m in json.load(r).get("data", [])]
    except urllib.error.HTTPError as error:
        result["error"] = f"/v1/models returned HTTP {error.code}; is the model still loading?"
        return result
    except OSError as error:
        result["error"] = f"cannot reach {base}: {error}"
        return result
    body = {
        "model": model, "stream": True, "max_tokens": max_tokens,
        "temperature": 1.0, "top_p": 1.0,
        "messages": [{"role": "user", "content": PROMPT}],
        "chat_template_kwargs": {"reasoning_effort": effort},
        "stream_options": {"include_usage": True},
    }
    content, reasoning, finish, usage, first = [], [], None, None, None
    start = time.monotonic()
    try:
        with request(f"{base}/v1/chat/completions", key, body, timeout) as r:
            for chunk in events(r):
                usage = chunk.get("usage") or usage
                for choice in chunk.get("choices", []):
                    delta = choice.get("delta", {})
                    text = delta.get("content") or ""
                    thought = delta.get("reasoning_content") or delta.get("reasoning") or ""
                    if (text or thought) and first is None:
                        first = time.monotonic() - start
                    content.append(text)
                    reasoning.append(thought)
                    finish = choice.get("finish_reason") or finish
    except urllib.error.HTTPError as error:
        result["error"] = f"chat request returned HTTP {error.code}: {error.read()[:300]!r}"
        return result
    except OSError as error:
        result["error"] = f"stream failed: {error}"
        return result
    answer = "".join(content).strip()
    result.update({
        "seconds_to_first_token": None if first is None else round(first, 2),
        "seconds_total": round(time.monotonic() - start, 2),
        "finish_reason": finish, "answer": answer,
        "reasoning_chars": len("".join(reasoning)), "usage": usage,
    })
    if finish == "length":
        result["error"] = "hit max_tokens before the answer finished; raise --max-tokens or lower --effort"
    elif finish != "stop":
        result["error"] = f"unexpected finish_reason {finish!r}"
    elif not answer:
        result["error"] = "the stream ended without answer text"
    else:
        result["ok"] = True
        result["answer_correct"] = "391" in answer
    return result


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--url", default="http://127.0.0.1:8080")
    p.add_argument("--model", default="muse-spark-open", help="the name /v1/models lists")
    p.add_argument("--effort", choices=EFFORTS, default="low")
    p.add_argument("--max-tokens", type=int, default=4096)
    p.add_argument("--timeout", type=float, default=600, help="seconds without data before giving up")
    a = p.parse_args(argv)
    result = run(a.url, a.model, a.effort, a.max_tokens, os.environ.get("MUSE_API_KEY"), a.timeout)
    print(json.dumps(result, indent=2))
    return 0 if result["ok"] else 1


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