refactor: store screenshots as URLs

report-app/src/app/pages/report-viewer/report-viewer.ts

@@ -236,10 +236,10 @@ export class ReportViewer {
 
   protected getScreenshotUrl(result: AssessmentResult): string | null {
     for (let i = result.attemptDetails.length - 1; i > -1; i--) {
-      const screenshot = result.attemptDetails[i].buildResult.screenshotBase64;
-
-      if (screenshot) {
-        return `data:image/png;base64,${screenshot}`;
+      const screenshotUrl =
+        result.attemptDetails[i].buildResult.screenshotPngUrl;
+      if (screenshotUrl) {
+        return screenshotUrl;
       }
     }
     return null;

runner/builder/builder-types.ts

@@ -56,7 +56,7 @@ export interface BuildResult {
   status: BuildResultStatus;
   message: string;
   errorType?: BuildErrorType;
-  screenshotBase64?: string; // Base64 encoded PNG screenshot
+  screenshotPngUrl?: string;
   missingDependency?: string;
   runtimeErrors?: string;
   /** JSON report from the Safety Web runner, if available. */

runner/builder/worker.ts

@@ -156,7 +156,7 @@ process.on('message', async (message: BuildWorkerMessage) => {
     result = {
       status: BuildResultStatus.SUCCESS,
       message: 'Application built successfully!',
-      screenshotBase64: screenshotBase64Data,
+      screenshotPngUrl: `data:image/png;base64,${screenshotBase64Data}`,
       runtimeErrors: runtimeErrors.join('\n'),
       axeViolations,
       safetyWebReportJson,

runner/file-system-utils.ts

@@ -65,7 +65,7 @@ export async function removeFolderWithSymlinks(dir: string) {
 /** Write a file and creates the necessary directory structure. */
 export async function safeWriteFile(
   filePath: string,
-  content: string,
+  content: string | Buffer,
   encoding?: BufferEncoding
 ): Promise<void> {
   const directory = dirname(filePath);

runner/ratings/autoraters/rate-files.ts

@@ -21,7 +21,7 @@ export async function autoRateFiles(
   environment: Environment,
   files: LlmResponseFile[],
   appPrompt: string,
-  screenshotBase64: string | null
+  screenshotPngUrl: string | null
 ): Promise<AutoraterRunInfo> {
   console.log(`Autorater is using '${model}' model. \n`);
 
@@ -39,15 +39,15 @@ export async function autoRateFiles(
 
   // Visual (screenshot) scoring...
   let visualRating = undefined;
-  if (screenshotBase64) {
+  if (screenshotPngUrl) {
     console.log('⏳ Awaiting visual scoring results...');
     visualRating = await autoRateAppearance(
       llm,
       abortSignal,
       model,
       environment,
       appPrompt,
-      screenshotBase64,
+      screenshotPngUrl,
       'command-line'
     );
     console.log(`${greenCheckmark()} Visual scoring is successful.`);

runner/ratings/autoraters/visuals-rater.ts

@@ -8,6 +8,7 @@ import {
 import { GenkitRunner } from '../../codegen/genkit/genkit-runner.js';
 import defaultVisualRaterPrompt from './visual-rating-prompt.js';
 import { Environment } from '../../configuration/environment.js';
+import { screenshotUrlToPngBuffer } from '../../utils/screenshots.js';
 
 /**
  * Automatically rate the appearance of a screenshot using an LLM.
@@ -16,7 +17,7 @@ import { Environment } from '../../configuration/environment.js';
  * @param model Model to use for the rating.
  * @param environment Environment in which the rating is running.
  * @param appPrompt Prompt to be used for the rating.
- * @param screenshotBase64 Screenshot to be rated.
+ * @param screenshotPngUrl Screenshot PNG URL to be rated.
  * @param label Label for the rating, used for logging.
  */
 export async function autoRateAppearance(
@@ -25,7 +26,7 @@ export async function autoRateAppearance(
   model: string,
   environment: Environment,
   appPrompt: string,
-  screenshotBase64: string,
+  screenshotPngUrl: string,
   label: string
 ): Promise<AutoRateResult> {
   const prompt = environment.renderPrompt(defaultVisualRaterPrompt, null, {
@@ -38,8 +39,10 @@ export async function autoRateAppearance(
       content: [
         {
           media: {
-            base64PngImage: screenshotBase64,
-            url: `data:image/png;base64,${screenshotBase64}`,
+            base64PngImage: (
+              await screenshotUrlToPngBuffer(screenshotPngUrl)
+            ).toString('base64'),
+            url: screenshotPngUrl,
           },
         },
       ],

runner/ratings/built-in-ratings/visual-appearance-rating.ts

@@ -20,7 +20,7 @@ export const visualAppearanceRating: LLMBasedRating = {
   id: 'common-autorater-visuals',
   model: DEFAULT_AUTORATER_MODEL_NAME,
   rate: async (ctx) => {
-    if (ctx.buildResult.screenshotBase64 === undefined) {
+    if (ctx.buildResult.screenshotPngUrl === undefined) {
       return {
         state: RatingState.SKIPPED,
         message: 'No screenshot available',
@@ -36,7 +36,7 @@ export const visualAppearanceRating: LLMBasedRating = {
         ctx.model,
         ctx.environment,
         ctx.fullPromptText,
-        ctx.buildResult.screenshotBase64,
+        ctx.buildResult.screenshotPngUrl,
         ctx.currentPromptDef.name
       );
     } catch (e) {

runner/reporting/report-logging.ts

@@ -105,13 +105,16 @@ export async function writeReportToDisk(
 
       // Write screenshot to fs first, since we'll remove this info
       // from JSON later in this function.
-      if (attempt.buildResult.screenshotBase64) {
+      if (attempt.buildResult.screenshotPngUrl) {
         const screenshotFilePath = join(attemptPath, 'screenshot.png');
-        await safeWriteFile(
-          screenshotFilePath,
-          attempt.buildResult.screenshotBase64,
-          'base64'
-        );
+
+        // Note: In practice this is a base64 data URL, but `fetch` conveniently
+        // allows us to extract the content for writing a PNG to disk.
+        const screenshotContent = await (
+          await fetch(attempt.buildResult.screenshotPngUrl)
+        ).arrayBuffer();
+
+        await safeWriteFile(screenshotFilePath, Buffer.from(screenshotContent));
       }
 
       // Write the safety web report if it exists.

runner/utils/screenshots.ts

@@ -0,0 +1,7 @@
+/** Converts a screenshot PNG URL to a PNG buffer with the image contents. */
+export async function screenshotUrlToPngBuffer(screenshotPngUrl: string) {
+  // Note: In practice this is a base64 data URL, but `fetch` conveniently
+  // allows us to extract the content for writing a PNG to disk.
+  const screenshotContent = await (await fetch(screenshotPngUrl)).arrayBuffer();
+  return Buffer.from(screenshotContent);
+}