Step 40: JSON Schema — schema-compiled constrained decoding (#78)

README.md

@@ -407,6 +407,7 @@ There is no NuGet package yet -- the project is in early development. Follow the
 
 ## News
 
+- **2026-03** — JSON Schema constrained decoding — `JsonSchemaConstraint` layers schema tracking on `JsonCharParser` to enforce type constraints, required properties, enum values, nested structures. Schema compiled into flat node array with property-name tries. Zero-alloc `Clone()` via struct-copy. `--response-format json_schema --schema <json|@file>` CLI support ([#78](https://github.com/kkokosa/dotLLM/issues/78))
 - **2026-03** — **Phase 5 begins**: JSON mode constrained decoding — `JsonConstraint` FSM guarantees syntactically valid JSON output via per-token vocabulary masking. Stack-based PDA (RFC 8259), AVX2-vectorized logit masking, state-keyed mask cache. `--response-format json_object` CLI flag ([#76](https://github.com/kkokosa/dotLLM/issues/76))
 - **2026-03** — KV-cache quantization: Q8_0 and Q4_0 KV-cache compression on CPU and GPU (3.7–7.1× memory reduction). Separate `--cache-type-k`/`--cache-type-v` options, mixed-precision window `--cache-window N` keeps recent tokens in full precision. Dual-region storage with quantize-on-evict, per-tile dequantization in tiled attention ([#74](https://github.com/kkokosa/dotLLM/issues/74))
 - **2026-03** — CPU/GPU hybrid layer offloading: `--gpu-layers N` to run first N layers on GPU, remainder on CPU. Automatic FP16→FP32 hidden state transfer at boundary. Split KV-cache (GPU FP16 + CPU FP32). Partial VRAM usage proportional to offloaded layers ([#72](https://github.com/kkokosa/dotLLM/issues/72))
@@ -440,7 +441,7 @@ There is no NuGet package yet -- the project is in early development. Follow the
 | **2 — Practical Local Inference** | Engine metrics, benchmarks, Q4_K_M, chat templates, streaming, multi-threading, more architectures | Done (10/10) |
 | **3 — CPU Performance** | Decode dispatch, Q8_1 input, weight repacking, outer-product GEMM, tiled attention, fast exp, fusion, NUMA | In Progress (7/8) |
 | **4 — GPU Acceleration** | CUDA backend, CPU/GPU hybrid, KV-cache quantization | Done (3/3) |
-| **5 — Constrained Decoding & API** | JSON mode, JSON Schema, regex/CFG, tool calling, logit bias, OpenAI API server | In Progress (1/6) |
+| **5 — Constrained Decoding & API** | JSON mode, JSON Schema, regex/CFG, tool calling, logit bias, OpenAI API server | In Progress (2/6) |
 | **6 — Production Serving** | Continuous batching, paged KV-cache, prompt caching, speculative decoding, metrics | Planned |
 | **7 — Expand** | Hooks, logit lens, LoRA, MLA, SAE, multi-GPU, ROCm | Planned |
 

docs/ROADMAP.md

@@ -93,7 +93,7 @@ Step 22 (done) ──────► Step 30 (NUMA + Spin-wait)
 | Step | Feature | Description | Depends On |
 |------|---------|-------------|------------|
 | 39 | **JSON mode** :white_check_mark: | `JsonConstraint` — FSM-based constrained decoding guaranteeing syntactically valid JSON. `response_format: {"type": "json_object"}`. | 8 |
-| 40 | **JSON Schema** | `JsonSchemaConstraint` — Schema-compiled automaton. `response_format: {"type": "json_schema", ...}`. Token mask precomputation. | 39 |
+| 40 | **JSON Schema** :white_check_mark: | `JsonSchemaConstraint` — Schema-compiled automaton. `response_format: {"type": "json_schema", ...}`. Token mask precomputation. | 39 |
 | 41 | **Regex + CFG** | `RegexConstraint` (DFA-based) and `GrammarConstraint` (PDA, GBNF-style). | 39 |
 | 42 | **Tool calling** | `IToolCallParser`, chat template tool integration, structured output for function arguments. `finish_reason: "tool_calls"`. Parallel tool calls. | 16, 40 |
 | 34 | **ASP.NET server** | Minimal API endpoints: `/v1/chat/completions`, `/v1/completions`, `/v1/models`, `/v1/embeddings`, `/v1/tokenize`, `/v1/detokenize`. Health + readiness probes. | 16, 17 |

scripts/test-schema.json

@@ -0,0 +1,9 @@
+{
+  "type": "object",
+  "properties": {
+    "name": { "type": "string" },
+    "age": { "type": "integer" }
+  },
+  "required": ["name", "age"],
+  "additionalProperties": false
+}

scripts/test_models_json_schema.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+"""
+test_models_json_schema.py — Smoke test for JSON Schema constrained decoding.
+
+Runs dotLLM CLI with --response-format json_schema and validates that the
+generated output conforms to the schema (correct types, required properties,
+no additional properties). Tests across cached model architectures.
+
+Usage:
+    python scripts/test_models_json_schema.py                    # run all cached models
+    python scripts/test_models_json_schema.py --filter llama     # only models matching "llama"
+    python scripts/test_models_json_schema.py --list             # show available test cases
+    python scripts/test_models_json_schema.py --device gpu       # GPU-only
+    python scripts/test_models_json_schema.py --device both      # CPU then GPU
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from bench_compare import resolve_model
+from test_models import TEST_CASES as MODEL_TEST_CASES, _model_is_cached as _base_model_is_cached
+
+
+# ---------------------------------------------------------------------------
+# Test configuration
+# ---------------------------------------------------------------------------
+
+SCHEMA_PATH = Path(__file__).resolve().parent / "test-schema.json"
+SCHEMA_PROMPT = "Output a JSON object with a person's name and age."
+SCHEMA_MAX_TOKENS = 60
+
+
+def _load_schema() -> dict:
+    """Load and parse the test schema."""
+    return json.loads(SCHEMA_PATH.read_text(encoding="utf-8"))
+
+
+def _find_cli() -> Path:
+    """Find the dotLLM CLI executable."""
+    repo_root = Path(__file__).resolve().parent.parent
+    bin_dir = repo_root / "src" / "DotLLM.Cli" / "bin"
+    for config in ("Release", "Debug"):
+        for ext in (".exe", ""):
+            p = bin_dir / config / "net10.0" / f"DotLLM.Cli{ext}"
+            if p.exists():
+                return p
+    return Path("dotnet")
+
+
+@dataclass
+class TestResult:
+    passed: bool
+    detail: str
+    elapsed: float
+    generated_text: str = ""
+
+
+def _validate_against_schema(parsed: dict, schema: dict) -> str | None:
+    """Validate parsed JSON against the test schema. Returns error string or None."""
+    props = schema.get("properties", {})
+    required = set(schema.get("required", []))
+    additional_forbidden = schema.get("additionalProperties") is False
+
+    # Check required properties
+    for key in required:
+        if key not in parsed:
+            return f"missing required property '{key}'"
+
+    # Check no additional properties
+    if additional_forbidden:
+        extra = set(parsed.keys()) - set(props.keys())
+        if extra:
+            return f"unexpected properties: {sorted(extra)}"
+
+    # Check types
+    type_map = {
+        "string": str,
+        "integer": int,
+        "number": (int, float),
+        "boolean": bool,
+        "null": type(None),
+    }
+    for key, prop_schema in props.items():
+        if key not in parsed:
+            continue
+        expected_type = prop_schema.get("type")
+        if expected_type and expected_type in type_map:
+            py_type = type_map[expected_type]
+            value = parsed[key]
+            if not isinstance(value, py_type):
+                # Special case: integer type should reject floats with fractional parts
+                if expected_type == "integer" and isinstance(value, float):
+                    if value != int(value):
+                        return f"'{key}' is float {value}, expected integer"
+                    # whole-number float is acceptable (JSON has no int type)
+                    continue
+                # bool is subclass of int in Python — reject bools for integer/number
+                if expected_type in ("integer", "number") and isinstance(value, bool):
+                    return f"'{key}' is bool, expected {expected_type}"
+                return f"'{key}' has type {type(value).__name__}, expected {expected_type}"
+
+    return None
+
+
+def _run_test(cli: Path, model_path: Path, schema: dict, tc, device: str = "cpu") -> TestResult:
+    """Run a single JSON Schema test."""
+    if cli.name == "dotnet":
+        cmd = [
+            str(cli), "run",
+            "--project", str(Path(__file__).resolve().parent.parent / "src" / "DotLLM.Cli"),
+            "-c", "Release", "--",
+        ]
+    else:
+        cmd = [str(cli)]
+
+    cmd += [
+        "run", str(model_path),
+        "-p", SCHEMA_PROMPT,
+        "-n", str(SCHEMA_MAX_TOKENS),
+        "-t", "0",
+        "--response-format", "json_schema",
+        "--schema", f"@{SCHEMA_PATH}",
+        "--json",
+        "--device", device,
+    ]
+
+    start = time.monotonic()
+    try:
+        result = subprocess.run(
+            cmd, capture_output=True, timeout=600,
+            encoding="utf-8", errors="replace",
+        )
+    except subprocess.TimeoutExpired:
+        return TestResult(False, "TIMEOUT (600s)", time.monotonic() - start)
+    except FileNotFoundError:
+        return TestResult(False, f"CLI not found: {cli}", time.monotonic() - start)
+
+    elapsed = time.monotonic() - start
+
+    if result.returncode != 0:
+        error_text = result.stderr.strip() or result.stdout.strip()
+        for line in error_text.splitlines():
+            if "Error:" in line:
+                return TestResult(False, line.strip(), elapsed)
+        return TestResult(False, f"exit code {result.returncode}: {error_text[:200]}", elapsed)
+
+    # Parse the CLI --json envelope
+    raw = result.stdout.strip()
+    json_start = raw.find("{")
+    if json_start < 0:
+        return TestResult(False, f"no JSON in CLI output: {raw[:200]}", elapsed)
+    try:
+        envelope = json.loads(raw[json_start:])
+    except json.JSONDecodeError as e:
+        return TestResult(False, f"invalid CLI JSON: {e}", elapsed)
+
+    generated_text = envelope.get("text", "").strip()
+
+    if not generated_text:
+        return TestResult(False, "empty output", elapsed, generated_text)
+
+    # Parse generated text as JSON
+    try:
+        parsed = json.loads(generated_text)
+    except json.JSONDecodeError as e:
+        preview = generated_text[:80].replace("\n", "\\n")
+        return TestResult(False, f"invalid JSON: {e} | {preview}", elapsed, generated_text)
+
+    if not isinstance(parsed, dict):
+        return TestResult(False, f"not object: {type(parsed).__name__}", elapsed, generated_text)
+
+    # Validate against schema
+    error = _validate_against_schema(parsed, schema)
+    if error:
+        preview = generated_text[:60].replace("\n", " ")
+        return TestResult(False, f"{error} | {preview}", elapsed, generated_text)
+
+    preview = generated_text[:60].replace("\n", " ")
+    return TestResult(True, preview, elapsed, generated_text)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Smoke test for JSON Schema constrained decoding."
+    )
+    parser.add_argument("--filter", type=str, default=None,
+                        help="Comma-separated name/arch substrings to match")
+    parser.add_argument("--list", action="store_true",
+                        help="List available test cases and exit")
+    parser.add_argument("--device", type=str, default="cpu",
+                        choices=["cpu", "gpu", "both"],
+                        help="Compute device (default: cpu)")
+    args = parser.parse_args()
+
+    cases = list(MODEL_TEST_CASES)
+    if args.filter:
+        filters = [f.strip().lower() for f in args.filter.split(",")]
+        cases = [
+            tc for tc in cases
+            if any(f in tc.name.lower() or f in tc.arch.lower() for f in filters)
+        ]
+        if not cases:
+            print(f"No test cases match filter '{args.filter}'.")
+            return 1
+
+    if args.list:
+        print(f"{'Name':<35} {'Arch':<10} {'Quant':<8} {'Cached':<8} Notes")
+        print("-" * 100)
+        for tc in cases:
+            cached = "yes" if _base_model_is_cached(tc) else "no"
+            quant = tc.quant or "default"
+            print(f"{tc.name:<35} {tc.arch:<10} {quant:<8} {cached:<8} {tc.notes}")
+        return 0
+
+    cases = [tc for tc in cases if _base_model_is_cached(tc)]
+    if not cases:
+        print("No cached models found. Download models first via test_models.py --download.")
+        return 1
+
+    schema = _load_schema()
+
+    cli = _find_cli()
+    if cli.name == "dotnet":
+        print("[cli] No prebuilt CLI found, will use 'dotnet run' (slower startup)")
+    else:
+        print(f"[cli] {cli}")
+
+    devices = ["cpu", "gpu"] if args.device == "both" else [args.device]
+    multi_device = len(devices) > 1
+
+    print(f"[schema] {SCHEMA_PATH}")
+    print(f"[schema] prompt: {SCHEMA_PROMPT}")
+    print(f"[schema] required: {schema.get('required', [])}, "
+          f"additionalProperties: {schema.get('additionalProperties', True)}")
+    print(f"[schema] max_tokens: {SCHEMA_MAX_TOKENS}")
+    print()
+    if multi_device:
+        print(f"{'Test':<30} {'Device':<6} {'Result':<8} {'Time':>8}  Details")
+    else:
+        print(f"{'Test':<30} {'Result':<8} {'Time':>8}  Details")
+    sep_width = 100 if multi_device else 90
+    print("=" * sep_width)
+
+    passed = 0
+    failed = 0
+
+    for tc in cases:
+        try:
+            model_path = resolve_model(tc.repo, tc.quant, quiet=True)
+        except SystemExit:
+            failed += 1
+            print(f"{tc.name:<30} {'FAIL':<8} {'':>8}  model resolution failed")
+            continue
+
+        for device in devices:
+            r = _run_test(cli, model_path, schema, tc, device=device)
+            time_str = f"{r.elapsed:.1f}s"
+            status = "PASS" if r.passed else "FAIL"
+            detail = r.detail[:55]
+
+            if r.passed:
+                passed += 1
+            else:
+                failed += 1
+
+            if multi_device:
+                print(f"{tc.name:<30} {device:<6} {status:<8} {time_str:>8}  {detail}")
+            else:
+                print(f"{tc.name:<30} {status:<8} {time_str:>8}  {detail}")
+
+    print("=" * sep_width)
+    total = passed + failed
+    print(f"\n{passed}/{total} passed, {failed} failed")
+    return 1 if failed > 0 else 0
+
+
+if __name__ == "__main__":
+    if sys.platform == "win32":
+        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+        sys.stderr.reconfigure(encoding="utf-8", errors="replace")
+    sys.exit(main())

src/DotLLM.Cli/Commands/ChatCommand.cs

@@ -123,10 +123,15 @@ public sealed class Settings : CommandSettings
 
         /// <summary>Constrain model output format.</summary>
         [CommandOption("--response-format")]
-        [Description("Constrain model output format: 'text' (default) or 'json_object' (valid JSON).")]
+        [Description("Constrain model output format: 'text' (default), 'json_object' (valid JSON), or 'json_schema' (schema-constrained JSON).")]
         [DefaultValue("text")]
         public string ResponseFormat { get; set; } = "text";
 
+        /// <summary>JSON Schema for json_schema response format.</summary>
+        [CommandOption("--schema")]
+        [Description("JSON Schema string or file path (prefixed with @) for json_schema response format.")]
+        public string? Schema { get; set; }
+
         /// <summary>KV-cache key quantization type.</summary>
         [CommandOption("--cache-type-k")]
         [Description("KV-cache key quantization: f32 (default), q8_0, q4_0.")]
@@ -220,9 +225,12 @@ public override async Task<int> ExecuteAsync(CommandContext context, Settings se
                 stopSequences.Add(marker);
         }
 
-        var responseFormat = settings.ResponseFormat.ToLowerInvariant() == "json_object"
-            ? (Core.Configuration.ResponseFormat)new Core.Configuration.ResponseFormat.JsonObject()
-            : null;
+        var responseFormat = settings.ResponseFormat.ToLowerInvariant() switch
+        {
+            "json_object" => (Core.Configuration.ResponseFormat)new Core.Configuration.ResponseFormat.JsonObject(),
+            "json_schema" => BuildJsonSchemaFormat(settings.Schema),
+            _ => null
+        };
         var inferenceOptions = new InferenceOptions
         {
             Temperature = settings.Temperature,
@@ -416,6 +424,18 @@ private static string InferQuantLabel(string resolvedPath, string? quantFlag)
         return match.Success ? match.Groups[1].Value : "unknown";
     }
 
+    private static Core.Configuration.ResponseFormat BuildJsonSchemaFormat(string? schema)
+    {
+        if (string.IsNullOrEmpty(schema))
+            throw new InvalidOperationException("--schema is required when --response-format is json_schema");
+
+        string schemaJson = schema.StartsWith('@')
+            ? File.ReadAllText(schema[1..])
+            : schema;
+
+        return new Core.Configuration.ResponseFormat.JsonSchema { Schema = schemaJson };
+    }
+
     private static string BuildSamplingLabel(Settings settings)
     {
         if (settings.Temperature <= 0)