Add JSON skeleton to compact ZeroGPU prompt

figment/model_client.py

@@ -15,7 +15,7 @@ import wave
 from typing import Any
 
 from .config import FigmentConfig, load_config
-from .prompt_builder import OUTPUT_SCHEMA
+from .prompt_builder import OUTPUT_SCHEMA, REQUIRED_JSON_SKELETON
 from .zerogpu_runtime import generate_zero_gpu_json
 
 
@@ -129,12 +129,15 @@ class ModelClient:
                 auth_headers={},
             )
         if self.config.model_backend == "hf_zerogpu":
-            return generate_zero_gpu_json(
-                prompt=_zero_gpu_compact_prompt(prompt, context),
-                model_repo=self.config.zerogpu_model_repo,
-                model_subfolder=self.config.zerogpu_model_subfolder,
-                model_id=self.config.local_model_id,
-            )
+            try:
+                return generate_zero_gpu_json(
+                    prompt=_zero_gpu_compact_prompt(prompt, context),
+                    model_repo=self.config.zerogpu_model_repo,
+                    model_subfolder=self.config.zerogpu_model_subfolder,
+                    model_id=self.config.local_model_id,
+                )
+            except (json.JSONDecodeError, RuntimeError, ValueError) as exc:
+                raise ModelClientError(f"hf_zerogpu backend failed for {self.config.local_model_id}") from exc
         if self.config.model_backend == "hosted_omni":
             endpoint = self.config.omni_endpoint_url or self.config.hf_endpoint_url or self.config.nvidia_base_url
             if not endpoint:
@@ -251,7 +254,8 @@ class ModelClient:
 
 
 ZERO_GPU_SYSTEM_PROMPT = """You are Figment, an offline protocol navigator for a trained responder.
-Use ONLY the compact protocol-card facts in CONTEXT. Return ONLY JSON with every required key.
+Use ONLY the compact protocol-card facts in CONTEXT. Return exactly one valid JSON object, with no markdown, prose, or comments.
+The object must follow REQUIRED_JSON_SKELETON and include every required key.
 Do not diagnose, prescribe, dose medication, autonomously triage, discharge, or downgrade deterministic red flags.
 Every candidate_protocol_pathways item must cite a retrieved clinical card_id. source_cards must cite every card used.
 If information is missing or uncertain, list it instead of inventing it."""
@@ -271,6 +275,7 @@ def _zero_gpu_compact_prompt(prompt: str, context: dict[str, Any]) -> str:
     }
     return (
         f"{ZERO_GPU_SYSTEM_PROMPT}\n\n"
+        f"REQUIRED_JSON_SKELETON:\n{json.dumps(REQUIRED_JSON_SKELETON, sort_keys=True, separators=(',', ':'))}\n\n"
         f"CONTEXT:\n{json.dumps(compact_context, sort_keys=True, separators=(',', ':'))}"
     )
 

tests/test_model_client.py

@@ -203,11 +203,29 @@ def test_hf_zerogpu_compact_prompt_preserves_route_facts_without_verbose_card_te
     assert len(compact) < 3000
     assert "PED-DEHYD-RED-FLAGS-v1" in compact
     assert "chief_concern" in compact
+    assert "REQUIRED_JSON_SKELETON" in compact
     assert "required_keys" in compact
     assert "source_note" not in compact
     assert "VERBOSE" not in compact
 
 
+def test_hf_zerogpu_backend_wraps_invalid_model_json(monkeypatch: Any) -> None:
+    def fake_generate_zero_gpu_json(**_: Any) -> dict[str, Any]:
+        raise json.JSONDecodeError("bad", "not json", 0)
+
+    monkeypatch.setattr("figment.model_client.generate_zero_gpu_json", fake_generate_zero_gpu_json)
+    client = ModelClient(
+        FigmentConfig(
+            model_backend="hf_zerogpu",
+            model_stack="local_4b_parakeet",
+            local_model_id="figment-sft-v14p-lora-merged-bf16",
+        )
+    )
+
+    with pytest.raises(ModelClientError, match="hf_zerogpu backend failed"):
+        client.generate_json("Return JSON.", {"intake": {}, "retrieved_cards": []})
+
+
 def test_custom_hf_endpoint_prefers_hf_token_over_nvidia_key(monkeypatch: Any) -> None:
     captured: dict[str, Any] = {}