fix: ignore truncated Gemini responses

Occasionally Gemini truncates the response on the `class=` sequence. These changes ignore such responses to reduce noise from the results.

Author: crisbeto

runner/codegen/genkit/genkit-runner.ts

@@ -45,7 +45,8 @@ export class GenkitRunner implements LlmRunner {
   async generateConstrained<T extends z.ZodTypeAny = z.ZodTypeAny>(
     options: LlmConstrainedOutputGenerateRequestOptions<T>
   ): Promise<LlmConstrainedOutputGenerateResponse<T>> {
-    const result = await this._genkitRequest(options);
+    const { provider, model } = this.resolveModel(options.model);
+    const result = await this._genkitRequest(provider, model, options);
 
     return {
       output: result.output,
@@ -72,10 +73,16 @@ export class GenkitRunner implements LlmRunner {
       }),
     };
 
-    const result = await this._genkitRequest(requestOptions);
+    const { provider, model } = this.resolveModel(options.model);
+    const result = await this._genkitRequest(provider, model, requestOptions);
+    const files = result.output.outputFiles || [];
+
+    if (!provider.validateGeneratedFiles(files)) {
+      throw new Error(`Invalid files generated by model "${options.model}"`);
+    }
 
     return {
-      files: result.output.outputFiles || [],
+      files,
       usage: result.usage,
       reasoning: result.reasoning,
       toolLogs: this.flushToolLogs(),
@@ -89,7 +96,8 @@ export class GenkitRunner implements LlmRunner {
   async generateText(
     options: LlmGenerateTextRequestOptions
   ): Promise<LlmGenerateTextResponse> {
-    const result = await this._genkitRequest(options);
+    const { provider, model } = this.resolveModel(options.model);
+    const result = await this._genkitRequest(provider, model, options);
 
     return {
       text: result.text,
@@ -104,12 +112,12 @@ export class GenkitRunner implements LlmRunner {
   }
 
   private async _genkitRequest(
+    provider: GenkitModelProvider,
+    model: ModelReference<any>,
     options:
       | LlmGenerateTextRequestOptions
       | LlmConstrainedOutputGenerateRequestOptions
   ) {
-    const { provider, model } = this.resolveModel(options.model);
-
     return await rateLimitLLMRequest(
       provider,
       model,

runner/codegen/genkit/model-provider.ts

@@ -2,6 +2,7 @@ import { ModelReference } from 'genkit';
 import { GenkitPlugin, GenkitPluginV2 } from 'genkit/plugin';
 import { RateLimiter } from 'limiter';
 import { PromptDataMessage } from '../llm-runner.js';
+import { LlmResponseFile } from '../../shared-interfaces.js';
 
 export interface RateLimitConfig {
   requestPerMinute: RateLimiter;
@@ -46,6 +47,12 @@ export abstract class GenkitModelProvider {
     return key ? this.pluginFactory(key) : null;
   }
 
+  /**
+   * Checks whether files generated by the LLM are valid.
+   * If not, the Genkit runner will throw an error.
+   */
+  abstract validateGeneratedFiles(files: LlmResponseFile[]): boolean;
+
   protected abstract pluginFactory(
     apiKey: string
   ): GenkitPlugin | GenkitPluginV2;

runner/codegen/genkit/providers/claude.ts

@@ -38,6 +38,10 @@ export class ClaudeModelProvider extends GenkitModelProvider {
     return {};
   }
 
+  validateGeneratedFiles(): boolean {
+    return true;
+  }
+
   private anthropicApi = lazy(() => {
     return new Anthropic({ apiKey: this.getApiKey() || undefined });
   });

runner/codegen/genkit/providers/gemini.ts

@@ -8,6 +8,7 @@ import {
 import { lazy } from '../../../utils/lazy-creation.js';
 import { GoogleGenAI, Part } from '@google/genai';
 import { RateLimiter } from 'limiter';
+import { LlmResponseFile } from '../../../shared-interfaces.js';
 
 export class GeminiModelProvider extends GenkitModelProvider {
   readonly apiKeyVariableName = 'GEMINI_API_KEY';
@@ -74,6 +75,15 @@ export class GeminiModelProvider extends GenkitModelProvider {
     return { thinkingConfig: { includeThoughts: opts.includeThoughts } };
   }
 
+  validateGeneratedFiles(files: LlmResponseFile[]): boolean {
+    // Gemini responses occasionally get truncated on `class=`.
+    // Consider these cases as invalid so they don't influence the results.
+    return (
+      files.length === 0 ||
+      !files.some((file) => file.code.trim().endsWith('class='))
+    );
+  }
+
   private async countGeminiTokens(
     prompt: PromptDataForCounting,
     modelName: string

runner/codegen/genkit/providers/grok.ts

@@ -96,6 +96,10 @@ export class GrokModelProvider extends GenkitModelProvider {
     return {};
   }
 
+  validateGeneratedFiles(): boolean {
+    return true;
+  }
+
   private genkitPromptToXaiFormat(
     prompt: PromptDataForCounting
   ): Array<{ role: string; content: string }> {

runner/codegen/genkit/providers/open-ai.ts

@@ -81,6 +81,10 @@ export class OpenAiModelProvider extends GenkitModelProvider {
     return {};
   }
 
+  validateGeneratedFiles(): boolean {
+    return true;
+  }
+
   private genkitPromptToOpenAi(
     prompt: PromptDataForCounting
   ): Array<{ role: string; content: string }> {