Trusting Gemini Structured Output in Production — Schema Design, Double Validation, and Bounded Retries

// invoice-schema.ts
import { Type } from "@google/genai";
 
export const invoiceSchema = {
  type: Type.OBJECT,
  properties: {
    invoiceNumber: { type: Type.STRING, description: "Invoice number" },
    issueDate: {
      type: Type.STRING,
      description: "Issue date. Always a real date in YYYY-MM-DD format",
    },
    vendor: {
      type: Type.OBJECT,
      properties: {
        name: { type: Type.STRING },
        taxId: { type: Type.STRING, description: "Registration number; omit if absent" },
      },
      required: ["name"],
    },
    lineItems: {
      type: Type.ARRAY,
      items: {
        type: Type.OBJECT,
        properties: {
          description: { type: Type.STRING },
          quantity: { type: Type.NUMBER },
          unitPrice: { type: Type.NUMBER, description: "Unit price, pre-tax" },
          amount: { type: Type.NUMBER, description: "quantity x unitPrice (pre-tax)" },
        },
        propertyOrdering: ["description", "quantity", "unitPrice", "amount"],
        required: ["description", "quantity", "unitPrice", "amount"],
      },
    },
    grandTotal: { type: Type.NUMBER, description: "Total including tax" },
    paymentStatus: {
      type: Type.STRING,
      enum: ["paid", "pending", "overdue"],
    },
    currency: { type: Type.STRING, enum: ["JPY", "USD", "EUR"] },
  },
  propertyOrdering: [
    "invoiceNumber", "issueDate", "vendor",
    "lineItems", "grandTotal", "paymentStatus", "currency",
  ],
  required: ["invoiceNumber", "issueDate", "lineItems", "grandTotal", "paymentStatus", "currency"],
};

// invoice-validate.ts
import { z } from "zod";
 
const InvoiceZ = z.object({
  invoiceNumber: z.string().min(1),
  issueDate: z.string().refine(
    (s) => !Number.isNaN(Date.parse(s)) && /^\d{4}-\d{2}-\d{2}$/.test(s),
    { message: "Not a real YYYY-MM-DD date" },
  ),
  vendor: z.object({ name: z.string().min(1), taxId: z.string().optional() }),
  lineItems: z.array(z.object({
    description: z.string(),
    quantity: z.number().positive(),
    unitPrice: z.number().min(0),
    amount: z.number().min(0),
  })).min(1),
  grandTotal: z.number().min(0),
  paymentStatus: z.enum(["paid", "pending", "overdue"]),
  currency: z.enum(["JPY", "USD", "EUR"]),
});
 
export type Invoice = z.infer<typeof InvoiceZ>;
 
// Verify cross-field consistency that the schema cannot express
export function checkConsistency(inv: Invoice): string[] {
  const issues: string[] = [];
 
  const lineSum = inv.lineItems.reduce((s, it) => s + it.amount, 0);
  // JPY has a minimum unit of 1 yen, so rounding noise is rare. For multi-currency, keep per-currency tolerances
  const tolerance = inv.currency === "JPY" ? 1 : 0.01;
  if (Math.abs(lineSum - inv.grandTotal) > tolerance) {
    issues.push(`Line sum ${lineSum} does not match grand total ${inv.grandTotal}`);
  }
 
  for (const it of inv.lineItems) {
    const expected = it.quantity * it.unitPrice;
    if (Math.abs(expected - it.amount) > tolerance) {
      issues.push(`Subtotal for "${it.description}" does not equal quantity x unitPrice`);
    }
  }
  return issues;
}
 
export function validateInvoice(raw: unknown):
  | { ok: true; data: Invoice; warnings: string[] }
  | { ok: false; error: string } {
  const parsed = InvoiceZ.safeParse(raw);
  if (!parsed.success) {
    return { ok: false, error: parsed.error.issues.map((i) => i.message).join("; ") };
  }
  return { ok: true, data: parsed.data, warnings: checkConsistency(parsed.data) };
}

// safe-generate.ts
import { GoogleGenAI } from "@google/genai";
 
const ai = new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY });
 
type GenResult<T> =
  | { ok: true; data: T }
  | { ok: false; reason: "truncated" | "blocked" | "parse" | "empty"; detail: string };
 
export async function safeGenerate<T>(
  prompt: string,
  schema: object,
  validate: (raw: unknown) => T,
): Promise<GenResult<T>> {
  const res = await ai.models.generateContent({
    model: "gemini-3.5-flash",
    contents: prompt,
    config: {
      responseMimeType: "application/json",
      responseSchema: schema,
      maxOutputTokens: 8192,
    },
  });
 
  const finish = res.candidates?.[0]?.finishReason;
  // MAX_TOKENS signals "needs more budget." Retrying with the same budget is wasted
  if (finish === "MAX_TOKENS") {
    return { ok: false, reason: "truncated", detail: "Reached maxOutputTokens" };
  }
  if (finish === "SAFETY" || finish === "PROHIBITED_CONTENT") {
    return { ok: false, reason: "blocked", detail: `Safety filter: ${finish}` };
  }
 
  const text = res.text;
  if (!text) return { ok: false, reason: "empty", detail: "Empty body" };
 
  try {
    return { ok: true, data: validate(JSON.parse(text)) };
  } catch (e) {
    return { ok: false, reason: "parse", detail: e instanceof Error ? e.message : String(e) };
  }
}

// pipeline.ts
import { safeGenerate } from "./safe-generate";
 
async function processOne<T>(
  doc: { id: string; content: string },
  schema: object,
  validate: (raw: unknown) => T,
  maxRetries = 2,
): Promise<{ id: string; status: "ok" | "review" | "failed"; data?: T; note?: string }> {
  let attempt = 0;
  while (attempt <= maxRetries) {
    const r = await safeGenerate(`Extract from:\n\n${doc.content}`, schema, validate);
    if (r.ok) return { id: doc.id, status: "ok", data: r.data };
 
    // truncated and blocked won't fix on retry. Send to a human immediately
    if (r.reason === "truncated" || r.reason === "blocked") {
      return { id: doc.id, status: "review", note: r.detail };
    }
    // parse / empty: throw again with exponential backoff
    await new Promise((res) => setTimeout(res, 800 * 2 ** attempt));
    attempt++;
  }
  return { id: doc.id, status: "failed", note: "Retry limit reached" };
}
 
// A simple worker pool that caps concurrency
export async function runPipeline<T>(
  docs: { id: string; content: string }[],
  schema: object,
  validate: (raw: unknown) => T,
  concurrency = 3,
) {
  const results: Awaited<ReturnType<typeof processOne<T>>>[] = [];
  const queue = [...docs];
  const workers = Array.from({ length: concurrency }, async () => {
    let doc;
    while ((doc = queue.shift())) {
      results.push(await processOne(doc, schema, validate));
    }
  });
  await Promise.all(workers);
  return results;
}