Hold the sampler to the action schema, so a small model calls the tool instead of describing it

distinct_agent/cli.py

@@ -1422,6 +1422,16 @@ def main(
             args.llama_server or _resolve_llama_server(args),
             energy_meter=energy_meter,
         )
+        # SAY WHICH DEVICE WILL DO THE WORK, AND SAY THAT IT WAS MEASURED.
+        #
+        # A volunteer with a graphics card who sees it sitting idle assumes
+        # something is broken. On the laptop this was developed on the card
+        # really is the wrong device -- 0.50 tokens per second against 2.93 on
+        # the CPU, because a partly offloaded model pays a round trip per
+        # token -- and the first time a model is loaded the worker spends a
+        # few minutes finding that out. Both are worth saying out loud.
+        runner._notify = lambda message: print(f"  {message}", file=sys.stderr, flush=True)
+        print(f"Runtime: {runner.executable}", file=sys.stderr, flush=True)
         if not runner.available:
             print(
                 "llama-server was not found. Supply --llama-server, or --llama-cli "

distinct_agent/harness.py

@@ -147,6 +147,7 @@ class StructuredToolHarness:
         # number chosen when nobody was measuring. See :func:`prompt_budget`.
         window = int(getattr(model.manifest, "context_length", 4096))
         limits = fit_output_tokens(job.limits, window)
+        limits["response_schema"] = action_schema(manifest)
         budget = prompt_budget(window, int(limits["max_output_tokens"]))
         per_result = min(self.max_result_characters, max(400, budget // 3))
         exchanges: list[str] = []
@@ -178,11 +179,15 @@ class StructuredToolHarness:
             )
             model_usages.append(dict(inference.usage))
             action = _parse_action(inference.text, self.max_action_characters)
-            if action is None and makers and not events and not nudged:
-                # Prose on the first move, when a tool here makes files. The
-                # parser accepts prose and always will, but accepting it here
-                # is how eleven of fourteen benchmark failures ended: a good
-                # paragraph *about* the document, and no document.
+            answered_without_acting = action is None or action["type"] == "final"
+            if answered_without_acting and makers and not events and not nudged:
+                # Finished on the first move, when a tool here makes files.
+                # Two shapes of the same failure: prose, which the parser
+                # accepts and always will, and a well-formed ``final`` that a
+                # constrained sampler will happily produce. Either way the
+                # answer is a good paragraph *about* the document and no
+                # document, which is how eleven of fourteen benchmark
+                # workloads failed.
                 nudged = True
                 current_prompt = _nudge_prompt(head, makers)
                 continue
@@ -790,6 +795,43 @@ def _example_arguments(schema: Any) -> dict[str, Any]:
     return example
 
 
+def action_schema(manifest: tuple[Mapping[str, Any], ...]) -> dict[str, Any]:
+    """The shape every turn of the structured loop must take.
+
+    Handed to the runner, which hands it to llama.cpp, which holds the sampler
+    to it. That is the difference between an instruction the model may ignore
+    and a shape it cannot leave, and it is the fix for the failure that cost
+    eleven of fourteen benchmark workloads: a model that wrote a paragraph
+    about the document instead of calling the tool that makes one.
+
+    Two alternatives rather than one loose object with optional fields: a
+    ``tool`` action without a tool, or a ``final`` without an answer, is
+    exactly the malformed action the parser has to refuse, and a schema that
+    permits it has not constrained anything worth constraining.
+    """
+
+    refs = [
+        str(item.get("ref") or f"{item.get('id')}@{item.get('version')}") for item in manifest
+    ]
+    call: dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "type": {"const": "tool"},
+            "tool": {"enum": refs} if refs else {"type": "string"},
+            "arguments": {"type": "object"},
+        },
+        "required": ["type", "tool", "arguments"],
+        "additionalProperties": False,
+    }
+    answer: dict[str, Any] = {
+        "type": "object",
+        "properties": {"type": {"const": "final"}, "answer": {"type": "string"}},
+        "required": ["type", "answer"],
+        "additionalProperties": False,
+    }
+    return {"anyOf": [call, answer]} if refs else answer
+
+
 def file_making_tools(manifest: tuple[Mapping[str, Any], ...]) -> tuple[str, ...]:
     """Which of these tools produce a file, read off their own schemas.
 

distinct_agent/server_runner.py

@@ -265,6 +265,10 @@ class LlamaServerRunner:
         #: ``None`` until something is loaded; ``0`` is a real answer meaning
         #: everything is on the CPU, and is not the same as "not measured".
         self.gpu_layers_used: int | None = None
+        #: Whether this build honours a JSON-schema constraint. ``None`` until
+        #: one has been tried; ``False`` after a build refuses one, which stops
+        #: every later request paying for the same refusal.
+        self.supports_schema: bool | None = None
         #: What calibration found, when it ran in this process. Empty when the
         #: answer came from the cache, which is the usual case.
         self.offload_measurements: tuple[Any, ...] = ()
@@ -336,6 +340,7 @@ class LlamaServerRunner:
         max_tokens: int,
         temperature: float,
         timeout: float,
+        schema: Mapping[str, Any] | None = None,
     ) -> tuple[Mapping[str, Any], str]:
         """One request lost its connection. Get the server back and ask again.
 
@@ -364,7 +369,11 @@ class LlamaServerRunner:
                 self._start_locked(model)
         try:
             return self._generate(
-                prompt, max_tokens=max_tokens, temperature=temperature, timeout=timeout
+                prompt,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                timeout=timeout,
+                schema=schema,
             )
         except _TransportLost as again:
             raise RunnerError(f"{again} (it had already been recovered once)") from again
@@ -641,6 +650,7 @@ class LlamaServerRunner:
         max_tokens: int,
         temperature: float,
         timeout: float,
+        schema: Mapping[str, Any] | None = None,
     ) -> tuple[Mapping[str, Any], str]:
         """Ask the model, through its own chat template wherever possible.
 
@@ -676,6 +686,22 @@ class LlamaServerRunner:
         omission.
         """
 
+        # ASKING NICELY FOR JSON DOES NOT WORK ON A SEVEN BILLION PARAMETER
+        # MODEL, AND IT DOES NOT HAVE TO.
+        #
+        # The harness needs one JSON object per turn. The prompt said so in
+        # capitals, gave a worked example, and re-asked once when prose came
+        # back; the model still wrote a perfectly good paragraph *about* the
+        # document it had been asked to create, called nothing, and produced
+        # no file. Eleven of fourteen benchmark failures were that.
+        #
+        # llama.cpp can constrain the sampler to a JSON schema, which turns
+        # "please reply in this shape" from an instruction the model may
+        # ignore into a shape it cannot leave. A build that does not support
+        # it says so with a 400, and the flag below stops it being asked
+        # again, so this degrades to the behaviour it replaces rather than
+        # failing.
+        constrain = schema is not None and self.supports_schema is not False
         chat_payload = {
             "messages": [{"role": "user", "content": prompt}],
             "max_tokens": max_tokens,
@@ -688,12 +714,33 @@ class LlamaServerRunner:
             # block that arrives anyway, because not every build honours this.
             "chat_template_kwargs": {"enable_thinking": False},
         }
+        if constrain:
+            chat_payload["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {"name": "action", "strict": True, "schema": dict(schema or {})},
+            }
         try:
-            return self._post(self.CHAT_PATH, chat_payload, timeout), self.PROMPT_MODE_CHAT
+            body = self._post(self.CHAT_PATH, chat_payload, timeout)
+            if constrain:
+                self.supports_schema = True
+            return body, self.PROMPT_MODE_CHAT
         except RunnerTimedOut:
             raise
         except _NoSuchEndpoint:
             pass
+        except RunnerError:
+            if not constrain:
+                raise
+            # Refused with a schema attached. Assume the schema is why, say so
+            # once, and carry on without it rather than losing the run.
+            self.supports_schema = False
+            return self._generate(
+                prompt,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                timeout=timeout,
+                schema=None,
+            )
         raw_payload = {
             "prompt": prompt,
             "n_predict": max_tokens,
@@ -701,6 +748,8 @@ class LlamaServerRunner:
             "stream": False,
             "cache_prompt": True,
         }
+        if constrain:
+            raw_payload["json_schema"] = dict(schema or {})
         return self._post(self.COMPLETION_PATH, raw_payload, timeout), self.PROMPT_MODE_RAW
 
     def _post(
@@ -819,12 +868,17 @@ class LlamaServerRunner:
 
         started = time.monotonic()
         temperature = float(limits.get("temperature", 0.2))
+        # Carried in limits rather than in the signature, so a runner that
+        # cannot constrain its output ignores it instead of refusing the call.
+        schema = limits.get("response_schema")
+        schema = schema if isinstance(schema, Mapping) else None
         try:
             body, prompt_mode = self._generate(
                 prompt,
                 max_tokens=max_tokens,
                 temperature=temperature,
                 timeout=timeout,
+                schema=schema,
             )
         except _TransportLost as lost:
             body, prompt_mode = self._retry_after_transport_loss(
@@ -834,6 +888,7 @@ class LlamaServerRunner:
                 max_tokens=max_tokens,
                 temperature=temperature,
                 timeout=timeout,
+                schema=schema,
             )
         except RunnerTimedOut:
             # LET IT FINISH TIDYING UP BEFORE THE NEXT RUN ARRIVES.
@@ -869,6 +924,10 @@ class LlamaServerRunner:
                 # different thing from a templated chat turn, and a run log has
                 # to be able to say which one it got.
                 "prompt_mode": prompt_mode,
+                # Whether the sampler was held to the action schema. An answer
+                # produced under a constraint is a different artefact from one
+                # produced freely, in the same way a templated turn is.
+                "schema_constrained": bool(schema) and self.supports_schema is not False,
                 "max_output_tokens": max_tokens,
                 "elapsed_seconds": round(time.monotonic() - started, 3),
                 "isolation": self.isolation.to_dict() if self.isolation else None,

tests/test_harness_makes_files.py

@@ -207,7 +207,9 @@ def test_reaching_the_budget_keeps_the_work_rather_than_discarding_it(tmp_path)
 
 
 def test_a_job_asking_for_more_than_the_worker_allows_is_clamped_not_refused(tmp_path) -> None:
-    broker = _Broker((DECK,), max_calls=2)
+    # No file-making tool here on purpose: this is about the budget, and with
+    # one present a first-turn final answer is correctly corrected instead.
+    broker = _Broker((CALCULATOR,), max_calls=2)
     result = _run(
         StructuredToolHarness(),
         _job(max_tool_calls=StructuredToolHarness.hard_max_tool_calls),
@@ -227,3 +229,75 @@ def test_a_job_asking_beyond_the_hard_ceiling_is_still_refused(tmp_path) -> None
             _Broker((DECK,)),
             tmp_path,
         )
+
+
+# -- the shape the sampler is held to ----------------------------------------
+
+
+def test_the_action_schema_offers_a_call_and_an_answer_and_nothing_else() -> None:
+    from distinct_agent.harness import action_schema
+
+    schema = action_schema((DECK, CALCULATOR))
+    shapes = schema["anyOf"]
+    assert [shape["properties"]["type"]["const"] for shape in shapes] == ["tool", "final"]
+    assert all(shape["additionalProperties"] is False for shape in shapes)
+
+
+def test_the_schema_only_permits_tools_that_exist() -> None:
+    """A constrained sampler cannot then invent a tool name."""
+
+    from distinct_agent.harness import action_schema
+
+    assert action_schema((DECK,))["anyOf"][0]["properties"]["tool"]["enum"] == ["create-deck@1"]
+
+
+def test_a_call_without_a_tool_is_not_a_permitted_shape() -> None:
+    """Constraining to something the parser would refuse constrains nothing."""
+
+    from distinct_agent.harness import action_schema
+
+    call = action_schema((DECK,))["anyOf"][0]
+    assert set(call["required"]) == {"type", "tool", "arguments"}
+
+
+def test_a_job_with_no_tools_can_only_answer() -> None:
+    from distinct_agent.harness import action_schema
+
+    schema = action_schema(())
+    assert "anyOf" not in schema
+    assert schema["properties"]["type"]["const"] == "final"
+
+
+def test_the_schema_travels_to_the_runner_with_the_run(tmp_path) -> None:
+    """It is carried in limits so a runner that cannot use it can ignore it."""
+
+    seen = {}
+
+    class _Watching(_Model):
+        def run(self, model, prompt, *, cancel_event=None, progress=None, limits=None):
+            seen.update(limits or {})
+            return super().run(model, prompt, cancel_event=cancel_event,
+                               progress=progress, limits=limits)
+
+    _run(
+        StructuredToolHarness(),
+        _job(),
+        _Watching(json.dumps({"type": "final", "answer": "done"})),
+        _Broker((CALCULATOR,)),
+        tmp_path,
+    )
+    assert "anyOf" in seen["response_schema"] or "properties" in seen["response_schema"]
+
+
+def test_a_well_formed_answer_that_skipped_the_file_is_corrected_too(tmp_path) -> None:
+    """A constrained sampler emits valid JSON. It can still answer without acting."""
+
+    runner = _Model(
+        json.dumps({"type": "final", "answer": "Here is an outline of the deck..."}),
+        json.dumps({"type": "tool", "tool": "create-deck@1", "arguments": {"slides": "A :: b"}}),
+        json.dumps({"type": "final", "answer": "deck built"}),
+    )
+    broker = _Broker((DECK,))
+    result = _run(StructuredToolHarness(), _job(), runner, broker, tmp_path)
+    assert broker.calls and broker.calls[0][0] == "create-deck@1"
+    assert result.text == "deck built"

tests/test_server_runner_recovery.py

@@ -293,3 +293,98 @@ def test_an_unreadable_error_body_is_simply_absent(tmp_path) -> None:
 
     error.read = explode
     assert server_runner._http_detail(error) == ""
+
+
+# -- holding the sampler to a shape ------------------------------------------
+
+
+SCHEMA = {
+    "anyOf": [
+        {
+            "type": "object",
+            "properties": {"type": {"const": "tool"}, "tool": {"enum": ["a@1"]}},
+            "required": ["type", "tool"],
+        },
+        {
+            "type": "object",
+            "properties": {"type": {"const": "final"}, "answer": {"type": "string"}},
+            "required": ["type", "answer"],
+        },
+    ]
+}
+
+
+def test_the_schema_reaches_llama_server_with_the_request(tmp_path, monkeypatch) -> None:
+    sent = {}
+
+    def post(self, path, payload, timeout):
+        sent["path"] = path
+        sent["payload"] = dict(payload)
+        return {"choices": [{"message": {"content": '{"type":"final","answer":"x"}'}}]}
+
+    monkeypatch.setattr(server_runner.LlamaServerRunner, "_post", post)
+    runner = _runner(tmp_path)
+    runner._generate("hi", max_tokens=10, temperature=0.0, timeout=5, schema=SCHEMA)
+    assert sent["payload"]["response_format"]["json_schema"]["schema"] == SCHEMA
+
+
+def test_a_build_that_refuses_the_schema_is_asked_again_without_it(tmp_path, monkeypatch) -> None:
+    """A constraint is an optimisation. Losing the run over it is not."""
+
+    attempts = []
+
+    def post(self, path, payload, timeout):
+        attempts.append("response_format" in payload)
+        if attempts[-1]:
+            raise RunnerError("llama-server rejected the request: HTTP 400")
+        return {"choices": [{"message": {"content": "plain text"}}]}
+
+    monkeypatch.setattr(server_runner.LlamaServerRunner, "_post", post)
+    runner = _runner(tmp_path)
+    body, _mode = runner._generate("hi", max_tokens=10, temperature=0.0, timeout=5, schema=SCHEMA)
+    assert attempts == [True, False]
+    assert runner.supports_schema is False
+
+
+def test_a_build_that_refused_once_is_not_asked_again(tmp_path, monkeypatch) -> None:
+    """Every later request would otherwise pay for the same refusal."""
+
+    attempts = []
+
+    def post(self, path, payload, timeout):
+        attempts.append("response_format" in payload)
+        return {"choices": [{"message": {"content": "text"}}]}
+
+    monkeypatch.setattr(server_runner.LlamaServerRunner, "_post", post)
+    runner = _runner(tmp_path)
+    runner.supports_schema = False
+    runner._generate("hi", max_tokens=10, temperature=0.0, timeout=5, schema=SCHEMA)
+    assert attempts == [False]
+
+
+def test_no_schema_means_no_constraint_in_the_payload(tmp_path, monkeypatch) -> None:
+    sent = {}
+
+    def post(self, path, payload, timeout):
+        sent.update(payload)
+        return {"choices": [{"message": {"content": "text"}}]}
+
+    monkeypatch.setattr(server_runner.LlamaServerRunner, "_post", post)
+    _runner(tmp_path)._generate("hi", max_tokens=10, temperature=0.0, timeout=5)
+    assert "response_format" not in sent
+
+
+def test_the_run_says_whether_its_answer_was_constrained(tmp_path, monkeypatch) -> None:
+    monkeypatch.setattr(
+        server_runner.LlamaServerRunner,
+        "_post",
+        lambda self, path, payload, timeout: {
+            "choices": [{"message": {"content": '{"type":"final","answer":"x"}'}}]
+        },
+    )
+    monkeypatch.setattr(server_runner.LlamaServerRunner, "ensure_started", lambda self, m: None)
+    runner = _runner(tmp_path)
+    result = runner.run(_model(tmp_path), "hi", limits={"response_schema": SCHEMA})
+    assert result.usage["schema_constrained"] is True
+    plain = runner.run(_model(tmp_path), "hi")
+    assert plain.usage["schema_constrained"] is False