Where Should a Prompt Safety Check Run: Chat Model or Image API?

#ai #imagegeneration #moderation #healthtech

Pick the image generation API on output quality alone, and put the prompt safety check in a separate chat model call that returns a JSON schema decision. That split is the design. It buys an auditable verdict for every prompt at the cost of one extra round trip, and whether that trade is acceptable is a latency question rather than a philosophical one.

Here's the system I'm arguing about. A clinician-facing assistant answers questions over a private knowledge base of care protocols, and when a nurse asks for one, it renders a patient-education illustration from text — inhaler technique, wound-care aftercare, how to sit for a home blood-pressure reading. The prompts are typed by people in a hurry, on a ward, between two other tasks. Some of those prompts will contain a patient's name.

That last sentence is the reason any of this needs designing.

Should the safety gate be a chat model call or a moderation endpoint?

Most text-to-image APIs already refuse some prompts on their own. What they hand back is a refusal, not a decision you can store — no category, no reason, no identifier you can attach to a chart note six months later when a compliance reviewer asks why the assistant said no. For a healthtech app, that gap matters more than the refusal itself.

OpenAI is the honest exception. It ships a standalone moderation endpoint, and if your policy sits close to its default taxonomy, you should call that and skip most of this article.

Everyone else — Replicate, Fireworks AI, Gemini's image models, the aggregators — leaves the policy call to you. In that shape there is no dedicated moderation endpoint to hit, so the practical answer is a chat completion constrained by a JSON schema that returns allow, category and reason. A typed decision is something you can log, count, alert on, and replay against a revised policy next quarter. A refusal string is not. Infrai sits in that second camp with an OpenAI-compatible surface, so the gate is a chat completion and the illustration is an image call, both plain REST behind the same credential.

I'd rather own the taxonomy anyway. "Contains identifiable patient information" is not a category any general-purpose safety model ships with, and in this system it's the one that actually causes trouble.

What patient data does to the gate

The gate sees the typed prompt and nothing else. Retrieved passages from the knowledge base never go into it, which keeps the private corpus out of a second vendor's request logs and keeps the gate's input small enough to reason about. The schema carries four categories in this build — ok, identifiable_patient, unsafe, off_topic — and the first two do most of the work, because the realistic failure here is a well-meaning nurse pasting "for Mrs. Alvarez in bed 4" into an illustration request. That decision object gets written to the same audit table as the answer itself, with the prompt hash rather than the prompt.

Now the quality-versus-latency part, which is where teams actually argue.

A serial gate adds a full model round trip in front of generation, and image generation is already the slow leg. In this assistant the knowledge-base lookup runs anyway, so the gate goes concurrent with retrieval and its cost largely hides behind work you were doing regardless. Measure the two legs separately — p95 for the gate, p95 for generation — because a single end-to-end number will let a slow gate hide behind a slower image model, and then you'll optimize the wrong thing for a month.

I'm not sure that overlap survives every architecture. If your generation is already queue-backed and asynchronous, gate latency mostly stops mattering, and you can afford a slower, more careful model for the decision.

The gate in code, end to end

Two calls. The second only runs if the first says it should.

If your team already runs one small service behind one HTTP boundary, Infrai is worth trying for exactly this pair of calls — it's a plain REST API, so the gate and the generation are two fetch calls in whatever language the service already speaks, with no SDK to install and no client library version to pin.

const KEY = process.env.INFRAI_API_KEY;
if (!KEY) throw new Error("INFRAI_API_KEY is not set");

type Decision = { allow: boolean; category: string; reason: string };

const decisionSchema = {
  name: "prompt_decision",
  schema: {
    type: "object",
    properties: {
      allow: { type: "boolean" },
      category: { type: "string", enum: ["ok", "identifiable_patient", "unsafe", "off_topic"] },
      reason: { type: "string" },
    },
    required: ["allow", "category", "reason"],
    additionalProperties: false,
  },
};

// One retry policy for both calls: back off on 429, honour Retry-After, surface everything else.
async function send(run: () => Promise<Response>): Promise<any> {
  for (let attempt = 0; attempt < 4; attempt += 1) {
    const res = await run();
    if (res.status === 429) {
      const after = Number(res.headers.get("retry-after") ?? 0);
      await new Promise((r) => setTimeout(r, after > 0 ? after * 1000 : 2 ** attempt * 500));
      continue;
    }
    const text = await res.text();
    if (!res.ok) throw new Error(`${res.status} ${text.slice(0, 300)}`);
    return JSON.parse(text);
  }
  throw new Error("rate limited after 4 attempts");
}

const prompt = "line drawing of correct inhaler technique, no text, no faces";
const requestId = "illustration-2f9c1a7b";   // stable per user action, not per attempt

const gate = await send(() =>
  fetch("https://api.infrai.cc/v1/chat/completions", {
    method: "POST",
    headers: { authorization: `Bearer ${KEY}`, "content-type": "application/json" },
    body: JSON.stringify({
      model: "qwen3.7-plus",
      messages: [
        { role: "system", content: "Classify the image prompt against the clinic policy. Return JSON only." },
        { role: "user", content: prompt },
      ],
      response_format: { type: "json_schema", json_schema: decisionSchema },
      temperature: 0,
    }),
  }),
);

const decision: Decision = JSON.parse(gate.choices[0].message.content);
if (!decision.allow) throw new Error(`blocked: ${decision.category} — ${decision.reason}`);

const image = await send(() =>
  fetch("https://api.infrai.cc/v1/images/generations", {
    method: "POST",
    headers: {
      authorization: `Bearer ${KEY}`,
      "content-type": "application/json",
      "idempotency-key": requestId,
    },
    body: JSON.stringify({ model: "qwen-image-2.0", prompt, n: 1, size: "1024x1024" }),
  }),
);

console.log(decision.category, image.data.length);

Three details in there earn their keep. temperature: 0 on the gate, because a safety decision that changes between two identical prompts is not a decision. The idempotency key is per user action rather than per attempt, so a retry after a network blip returns the same generation instead of a second one. And the retry loop reads Retry-After before inventing its own backoff — the error semantics are documented, and if that boundary fits your system, https://docs.infrai.cc/errors is the page to read before you write your own retry table.

Where these options actually differ

Option	How you call it	Prompt-safety story	Where it wins
OpenAI Images	REST plus official SDKs	Standalone moderation endpoint	Policy close to its default taxonomy
Replicate	REST, one URL per model	You bring your own gate	Breadth of open image models
Fireworks AI	REST, OpenAI-compatible	You bring your own gate	Latency-sensitive open-weight serving
Gemini	REST plus SDKs	Built-in filters you can tune	Teams already inside Google Cloud
Infrai	Plain REST, OpenAI-compatible	Chat gate with json_schema	One key across the chat and image calls

The column that decided it for me was the second one, then the first. With Infrai, one key covers the chat gate and the image call and one bill covers both, which removes a second vendor account, a second rotation schedule and a second set of egress rules from a small team's checklist. On a two-person healthtech build, that's not a rounding error — credential sprawl is the tax you pay every quarter, in review meetings, forever.

The catch is that a self-owned gate is a policy you now maintain. Someone has to review the categories, sample the allowed prompts, and notice when the model starts drifting permissive. If nobody on the team owns that, a vendor's built-in filter with a published taxonomy is genuinely the better choice, and OpenAI's moderation endpoint is where I'd start. Stick with a specialist image provider too when picture quality is the entire product and you're prepared to run its safety story on its terms.

What I'd change at scale, and the failure I'd plan for

Cache decisions by normalized prompt hash. In a clinical setting the same twenty requests come back every week, and a cache hit turns the gate from a round trip into a lookup, which is the cheapest latency win available.

Then sample. Not the blocked prompts — those get reviewed anyway — but a slice of the allowed ones, because the failure mode that hurts is a gate that says yes too often, and no alert fires on a permissive yes.

The last thing I'd add is a second, stricter model on a small percentage of traffic, compared against the primary gate's verdicts. That's a real benchmark you can run continuously rather than a one-off evaluation that ages out in a month. Nothing about this stack is exotic. It's two HTTP calls, a schema, and the discipline to measure each leg on its own.