From c2bfe33a6ad1e4f9fa3cf14b49d6156c3c7dee76 Mon Sep 17 00:00:00 2001 From: anandpant <109482096+anandpant@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:25:46 +0000 Subject: [PATCH] fix(generation): expand output budget (#310) ## Scope - Raise Gemini's default output budget from 2,048 to 16,384 tokens while preserving request overrides. - Train compact one-line JSON through both instructions and examples. - Classify Gemini `MAX_TOKENS` completion as explicit `output_truncated` failure evidence. ## Proof - `pnpm nx run-many -t typecheck,test,build` - passed (52 tasks; 50 executed, 2 cache hits) - `pnpm nx build-storybook diagram-ui` - passed - `pnpm nx test diagram-scenarios --skip-nx-cache` - passed - Added default-budget, compact-message, and truncation classification tests. This is PR 2 of 6 in the generation-reliability stack. --- .../src/lib/cloudflare-google-ai-studio.ts | 8 ++ packages/diagram/generation/src/lib/gemini.ts | 11 ++- .../diagram/generation/src/lib/messages.ts | 92 +++++++++---------- .../generation/src/lib/public-api.test.ts | 92 ++++++++++++++----- .../diagram/scenarios/src/lib/prompt.test.ts | 6 +- 5 files changed, 132 insertions(+), 77 deletions(-) diff --git a/packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts b/packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts index 319eefe8..635db551 100644 --- a/packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts +++ b/packages/diagram/generation/src/lib/cloudflare-google-ai-studio.ts @@ -15,6 +15,7 @@ import { } from "./errors.js"; import { buildGeminiGenerateContentBody, + extractGeminiFinishReason, extractGeminiText, extractGeminiUsage, stripGoogleModelPrefix, @@ -183,8 +184,15 @@ const runGatewayAttempt = Effect.fn( }), }); const usage = extractGeminiUsage(raw); + const finishReason = extractGeminiFinishReason(raw); return candidateFromText({ + diagnostics: + finishReason === "MAX_TOKENS" + ? [ + "output_truncated: Gemini stopped at the maximum output-token budget; regenerate the complete diagram.", + ] + : [], model, provider: "cloudflare-google-ai-studio", raw, diff --git a/packages/diagram/generation/src/lib/gemini.ts b/packages/diagram/generation/src/lib/gemini.ts index 20df09a2..10904a51 100644 --- a/packages/diagram/generation/src/lib/gemini.ts +++ b/packages/diagram/generation/src/lib/gemini.ts @@ -36,7 +36,7 @@ export class GeminiGenerateContentBody extends Schema.Class 0 ? usage : undefined; } + +export function extractGeminiFinishReason( + response: unknown, +): string | undefined { + const candidates = objectValue(response, "candidates"); + const firstCandidate = Array.isArray(candidates) ? candidates[0] : undefined; + const finishReason = objectValue(firstCandidate, "finishReason"); + return typeof finishReason === "string" ? finishReason : undefined; +} diff --git a/packages/diagram/generation/src/lib/messages.ts b/packages/diagram/generation/src/lib/messages.ts index ac46cc41..cc048461 100644 --- a/packages/diagram/generation/src/lib/messages.ts +++ b/packages/diagram/generation/src/lib/messages.ts @@ -36,7 +36,7 @@ export class DiagramGenerationMessages extends Schema.Class { expect(messages.system).toContain('Use type "mindmap".'); expect(messages.system).toContain("exactly one root node"); - expect(messages.user).toContain('"type": "mindmap"'); + expect(messages.user).toContain('"type":"mindmap"'); }); it("keeps the Gemini REST body byte-for-byte stable", () => { @@ -216,6 +212,15 @@ describe("diagram generation prompt mapping", () => { ).toEqual(expectedGeminiBody); }); + it("uses the expanded output budget unless the request overrides it", () => { + expect( + buildGeminiGenerateContentBody({ + model: "google/gemini-3.1-flash-lite", + prompt, + }).generationConfig.maxOutputTokens, + ).toBe(16_384); + }); + it("normalizes Cloudflare Google model ids for provider-native calls", () => { expect(stripGoogleModelPrefix("google/gemini-3.1-flash-lite")).toBe( "gemini-3.1-flash-lite", @@ -774,6 +779,45 @@ layer(malformedClientLayer)("malformed model responses", (it) => { ); }); +const truncatedRun = vi.fn(async () => + jsonResponse({ + candidates: [ + { + content: { role: "model", parts: [{ text: '{"type":"flowchart"' }] }, + finishReason: "MAX_TOKENS", + }, + ], + }), +); +const truncatedClientLayer = CloudflareGoogleAiStudioClientLive.pipe( + Layer.provide( + Layer.mergeAll( + Layer.succeed(CloudflareAiGatewayBinding, { + gateway: () => ({ getUrl: vi.fn(), run: truncatedRun }), + }), + configLayer, + retryPolicyLayer, + ), + ), +); + +layer(truncatedClientLayer)("token-budget exhaustion", (it) => { + it.effect("returns an explicit output-truncated diagnostic", () => + Effect.gen(function* () { + const client = yield* DiagramGenerationClient; + const candidate = yield* client.generate({ + model: "google/gemini-3.1-flash-lite", + prompt, + }); + + assert.isUndefined(candidate.diagram); + expect(candidate.diagnostics).toContain( + "output_truncated: Gemini stopped at the maximum output-token budget; regenerate the complete diagram.", + ); + }), + ); +}); + const substitutedClientLayer = Layer.succeed(DiagramGenerationClient, { provider: "fixture", generate: Effect.fn("diagramGeneration.test.generate")(function* () { diff --git a/packages/diagram/scenarios/src/lib/prompt.test.ts b/packages/diagram/scenarios/src/lib/prompt.test.ts index 79fc2704..040a0721 100644 --- a/packages/diagram/scenarios/src/lib/prompt.test.ts +++ b/packages/diagram/scenarios/src/lib/prompt.test.ts @@ -45,7 +45,7 @@ describe("scenario prompts", () => { expect(prompt.user).toContain("- QA Manager final review"); expect(prompt.user).toContain("Required decision branch labels:"); expect(prompt.user).toContain("- retest"); - expect(prompt.user).toContain('"title": "Pharma batch disposition"'); + expect(prompt.user).toContain('"title":"Pharma batch disposition"'); }); it("keeps a flattened prompt for stdin-based generator commands", () => { @@ -53,7 +53,9 @@ describe("scenario prompts", () => { expect(flattened).toContain("System message:"); expect(flattened).toContain("User message:"); - expect(flattened).toContain("Return only JSON."); + expect(flattened).toContain( + "Return only compact, minified JSON on one line.", + ); expect(flattened).toContain(scenario.prompt); }); });