feat: add function for augmenting generated files

Adds the ability for the user to augment a generated response before it is evaluated.

Author: crisbeto

runner/configuration/environment-config.ts

@@ -8,7 +8,12 @@ import {
   LocalExecutorConfig,
   localExecutorConfigSchema,
 } from '../orchestration/executors/local-executor-config.js';
-import {PromptDefinition, RatingContextFilter, ReportContextFilter} from '../shared-interfaces.js';
+import {
+  LlmResponseFile,
+  PromptDefinition,
+  RatingContextFilter,
+  ReportContextFilter,
+} from '../shared-interfaces.js';
 import type {Environment} from './environment.js';
 import type {GenkitRunner} from '../codegen/genkit/genkit-runner.js';
 
@@ -127,6 +132,13 @@ export const environmentConfigSchema = z.object({
   augmentExecutablePrompt: z
     .function(z.tuple([z.custom<PromptAugmentationContext>()]), z.promise(z.string()))
     .optional(),
+
+  /**
+   * Function that can be used to augment generated files before they're evaluated.
+   */
+  augmentGeneratedFile: z
+    .function(z.tuple([z.custom<Readonly<LlmResponseFile>>()]), z.string())
+    .optional(),
 });
 
 /**

runner/configuration/environment.ts

@@ -5,6 +5,7 @@ import {Executor} from '../orchestration/executors/executor.js';
 import {Rating, RatingCategory} from '../ratings/rating-types.js';
 import {
   FrameworkInfo,
+  LlmResponseFile,
   MultiStepPromptDefinition,
   PromptDefinition,
   RatingContextFilter,
@@ -83,6 +84,9 @@ export class Environment {
   /** Runner that user can use to access an LLM to augment prompts. */
   private augmentationRunner: GenkitRunner | null = null;
 
+  /** User-provided callback for augmenting the LLM-generated files. */
+  private readonly augmentFileCallback: ((file: LlmResponseFile) => string) | null;
+
   constructor(
     rootPath: string,
     private readonly config: EnvironmentConfig & Required<Pick<EnvironmentConfig, 'executor'>>,
@@ -114,6 +118,7 @@ export class Environment {
     this.ratingHash = this.getRatingHash(this.ratings, this.ratingCategories);
     this.analysisPrompts = this.resolveAnalysisPrompts(config);
     this.augmentExecutablePrompt = config.augmentExecutablePrompt || null;
+    this.augmentFileCallback = config.augmentGeneratedFile || null;
     this.validateRatingHash(this.ratingHash, config);
   }
 
@@ -191,6 +196,13 @@ export class Environment {
     });
   }
 
+  /** Augments response files based on the user's configuration. */
+  augmentResponseFiles(files: LlmResponseFile[]): void {
+    if (this.augmentFileCallback) {
+      files.forEach(file => (file.code = this.augmentFileCallback!(file)));
+    }
+  }
+
   async destroy(): Promise<void> {
     await this.executor.destroy();
 

runner/orchestration/codegen.ts

@@ -129,6 +129,7 @@ export async function repairCodeWithAI(
   );
 
   if (response.success) {
+    env.augmentResponseFiles(response.outputFiles);
     progress.log(
       promptDef,
       'codegen',

runner/orchestration/generate-initial-files.ts

@@ -75,6 +75,7 @@ export async function generateInitialFiles(
   );
 
   if (response.success) {
+    env.augmentResponseFiles(response.outputFiles);
     progress.log(
       promptDef,
       'codegen',
@@ -90,7 +91,7 @@ export async function generateInitialFiles(
   }
 
   return {
-    files: response.outputFiles!,
+    files: response.outputFiles,
     usage: response.usage,
     reasoning: response.reasoning,
     toolLogs: response.toolLogs,

runner/orchestration/repair.ts

@@ -126,7 +126,7 @@ async function handleRepairResponse(
   const newAttemptFiles = previousAttemptFiles.map(f => ({...f}));
 
   mergeRepairFiles(repairResponse.outputFiles, newAttemptFiles);
-  writeResponseFiles(directory, newAttemptFiles, env, rootPromptDef.name);
+  await writeResponseFiles(directory, newAttemptFiles, env, rootPromptDef.name);
 
   const buildResult = await runBuild(
     evalID,