# Constraining LLMs: Implementing Type-Safe Structured Outputs with Zod and JSON Schema

URL: https://mepritam.dev/references/constraining-llms-structured-outputs-zod-jsonschema/

A technical guide on bypassing the fragility of prompt engineering by implementing runtime-validated, schema-constrained LLM outputs in production TypeScript pipelines.

---

For a long time, integrating Large Language Models (LLMs) into traditional software systems felt like trying to splice a biological organ into a machine. You are connecting a probabilistic engine that outputs natural language to a deterministic codebase that expects typed parameters, strict interfaces, and predictable data structures.

Early strategies relied on "prompt engineering"—begging the model to _"please respond in valid JSON"_—followed by brittle regular expressions or `JSON.parse` wrappers wrapped in massive `try/catch` blocks. In production, this approach fails. A single stray character, an unescaped double quote, or a slight change in the model's tone breaks downstream services.

To build robust applications, we must move from prompt pleading to **schema-constrained generation**. This guide dives into how constrained outputs work under the hood and how to implement a type-safe extraction pipeline in TypeScript using Zod and JSON Schema.

---

## 1. How Constrained Generation Works Under the Hood

To understand why native structured outputs are more reliable than prompt-engineering hacks, we need to look at how token sampling works at the inference layer.

An LLM generates text token by token. At each step, it outputs a probability distribution over its entire vocabulary (tens of thousands of tokens). Normally, the sampling algorithm (like top-k, top-p, or temperature) picks a token from this distribution.

![LLM Structured Outputs Workflow](/images/references/structured-outputs-flow.jpg)

When you specify a **JSON Schema** constraint in modern APIs (such as OpenAI's Structured Outputs or Gemini's Schema-constrained JSON):

1. **Context-Free Grammar (CFG) Masking:** The engine constructs a parser or state machine based on the JSON Schema.
2. **Logit Bias / Masking:** At each token generation step, the engine identifies which tokens from the vocabulary would violate the grammar state. For example, if the model has just generated `{"age": `, the next valid tokens must be numbers or whitespace. A closing brace `}` or a string like `"young"` is syntactically invalid.
3. **Dynamic Filtering:** The probabilities of invalid tokens are set to `-Infinity` (masked out). The model is forced to sample _only_ from the subset of tokens that conform to the schema structure.

By the time the output reaches your server, it is mathematically guaranteed to be syntactically valid JSON matching your schema structure. This eliminates syntax parsing errors entirely.

---

## 2. Unstructured vs. Structured Output: The Developer's Perspective

The architectural benefits of type-safe outputs are clear when comparing the two retrieval styles:

![Structured vs. Unstructured LLM Output Comparison](/images/references/text-vs-structured-json.jpg)

With unstructured text, you inherit downstream complexity. You must write parsers, sanitizers, and write defensive code for every field. With structured JSON, the LLM functions exactly like a strongly-typed REST API.

---

## 3. Implementing Constrained Generation in TypeScript

Let’s write a production-ready module to extract structured product specifications from raw, unorganized vendor descriptions. We will use **Zod** to define our application-level types and the official SDKs to enforce the schema.

### Step 1: Define the Zod Schema

First, we define a schema for the structured data we want to extract. We want clean keys, proper types, and validation constraints.

```typescript
import { z } from "zod";

export const ProductSpecsSchema = z.object({
  productName: z.string().describe("The clean, standard name of the product."),
  brand: z
    .string()
    .nullable()
    .describe("The manufacturer/brand, or null if unknown."),
  categories: z.array(z.string()).describe("List of relevant category tags."),
  attributes: z.array(
    z.object({
      key: z
        .string()
        .describe("The attribute name, e.g., 'Weight', 'Battery Life'"),
      value: z
        .string()
        .describe("The value of the attribute, e.g., '1.2kg', '12 hours'"),
      unit: z
        .string()
        .optional()
        .describe("Measurement unit if applicable, e.g., 'kg', 'hrs'"),
    }),
  ),
  pricing: z
    .object({
      amount: z.number().positive().describe("Extracted price amount."),
      currency: z
        .string()
        .length(3)
        .default("USD")
        .describe("3-letter ISO currency code."),
    })
    .nullable()
    .describe("Extracted price, or null if no price is mentioned."),
});

export type ProductSpecs = z.infer<typeof ProductSpecsSchema>;
```

### Step 2: Query the LLM with Schema Constraints

Next, we pass this schema to the LLM. We will write an implementation using standard APIs. We transform the Zod schema into a JSON Schema format that the model's inference engine can understand.

```typescript
import { zodToJsonSchema } from "zod-to-json-schema";
import { GoogleGenAI } from "@google/genai"; // Example with Gemini SDK

const ai = new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY });

export async function extractProductSpecs(
  rawDescription: string,
): Promise<ProductSpecs> {
  // Convert Zod schema to raw JSON Schema format
  const jsonSchema = zodToJsonSchema(ProductSpecsSchema);

  const response = await ai.models.generateContent({
    model: "gemini-2.5-flash",
    contents: [
      {
        role: "user",
        parts: [
          {
            text: "Extract structural product specifications from the following text:",
          },
          { text: rawDescription },
        ],
      },
    ],
    config: {
      responseMimeType: "application/json",
      // Enforce the schema constraints at the inference layer
      responseSchema: jsonSchema as any,
    },
  });

  const rawText = response.text;
  if (!rawText) {
    throw new Error("Received empty response from the AI model.");
  }

  // Parse the guaranteed JSON syntax and pass it through Zod for semantic validation
  const parsedData = JSON.parse(rawText);
  return ProductSpecsSchema.parse(parsedData);
}
```

---

## 4. Handling Semantic Validation and the Auto-Retry Loop

While logit-masking guarantees _syntactic_ validity (the brackets and quotes match), it cannot guarantee _semantic_ validation. For example, if Zod expects a 3-letter currency code (like `USD` or `EUR`) via `.length(3)`, but the model outputs `USDT` or `DOLLARS`, the JSON parser succeeds, but Zod throws a validation error.

In a production system, you should not crash on validation failures. Instead, build an **automatic self-repair loop** that feeds the Zod validation error back to the model.

```typescript
export async function extractWithRetry(
  rawDescription: string,
  maxRetries = 3,
): Promise<ProductSpecs> {
  let attempts = 0;
  let feedbackPrompt = "";

  while (attempts < maxRetries) {
    attempts++;
    try {
      const response = await ai.models.generateContent({
        model: "gemini-2.5-flash",
        contents: [
          {
            role: "user",
            parts: [
              {
                text: "Extract structural specs from the raw description below.",
              },
              { text: `Raw Description:\n${rawDescription}` },
              feedbackPrompt
                ? {
                    text: `\n\n[System Alert: Previous Attempt Failed]\n${feedbackPrompt}`,
                  }
                : null,
            ].filter((p): p is { text: string } => p !== null),
          },
        ],
        config: {
          responseMimeType: "application/json",
          responseSchema: zodToJsonSchema(ProductSpecsSchema) as any,
        },
      });

      const parsed = JSON.parse(response.text || "{}");

      // Perform strict validation
      return ProductSpecsSchema.parse(parsed);
    } catch (error) {
      if (error instanceof z.ZodError) {
        console.warn(`Attempt ${attempts} failed validation:`, error.errors);

        // Format Zod errors into a helpful prompt for the next loop
        const errorDetails = error.errors
          .map(
            (e) =>
              `- Path "${e.path.join(".")}": ${e.message} (received invalid value)`,
          )
          .join("\n");

        feedbackPrompt = `Your previous JSON output did not meet the semantic rules. Please correct these errors:\n${errorDetails}`;
      } else {
        // Network errors or API rate limits should throw immediately
        throw error;
      }
    }
  }

  throw new Error(
    `Failed to extract valid specifications after ${maxRetries} attempts.`,
  );
}
```

---

## 5. Architectural Takeaways for Senior Developers

Adding schema constraint to your LLM pipelines changes your system architecture in several key ways:

- **Type Safety at the Boundary:** By linking your API validation (Zod) directly to your AI model configuration, you maintain a single source of truth for your data models.
- **Robust Integration Testing:** Because the output is predictable, you can write mock fixtures, perform unit tests on your transformation layers, and run Jest tests against mock outputs without dealing with string variability.
- **Improved Security (Prompt Injection Mitigation):** Attackers attempting prompt injection inside `rawDescription` find it much harder to escape the boundary. Since the model is constrained by logit biases, it physically cannot output malicious instructions outside the expected schema fields.

By treating LLMs as structured, typed components rather than black-box text generators, you can build reliable AI features that integrate into enterprise workflows.