Weights & Biases integration

package.json

@@ -62,6 +62,7 @@
   "dependencies": {
     "@anthropic-ai/sdk": "^0.5.2",
     "@apidevtools/json-schema-ref-parser": "^10.1.0",
+    "@wandb/sdk": "^0.5.1",
     "async": "^3.2.4",
     "cache-manager": "^4.1.0",
     "cache-manager-fs-hash": "^1.0.0",

src/main.ts

@@ -202,6 +202,16 @@ async function main() {
       'This suffix is append to every prompt',
       config.defaultTest?.options?.suffix,
     )
+    .option(
+      '--wandb-project',
+      'Enable logging to a Weights & Biases project. The project must already exist in the W&B dashboard.',
+      config?.commandLineOptions?.wandbProject,
+    )
+    .option(
+      '--wandb-name',
+      'Enable logging to a Weights & Biases under a run name',
+      config?.commandLineOptions?.wandbName,
+    )
     .option(
       '--no-write',
       'Do not write results to promptfoo directory',
@@ -303,6 +313,27 @@ async function main() {
         options.generateSuggestions = true;
       }
 
+      // Other test setup
+      let wandb;
+      if (cmdObj.wandbProject || cmdObj.wandbName) {
+        if (!process.env.WANDB_API_KEY) {
+          throw new Error(
+            'You must set WANDB_API_KEY in your environment to use the --wandb-project or --wandb-name options',
+          );
+        }
+        logger.info('Enabled logging to Weights & Biases');
+
+        wandb = (await import('@wandb/sdk')).default;
+
+        await wandb.init({
+          project: cmdObj.wandbProject,
+          name: cmdObj.wandbName,
+          config: testSuite,
+        });
+      }
+
+      // Run it!
+
       const summary = await evaluate(testSuite, options);
 
       const shareableUrl = cmdObj.share ? await createShareableUrl(summary, config) : null;
@@ -321,6 +352,26 @@ async function main() {
         }
       }
 
+      // Outputs
+
+      if (wandb) {
+        logger.info('Sending everything to Weights & Biases...');
+
+        // Record all results
+        for (const result of summary.results) {
+          wandb.log({
+            prompt: result.prompt.raw,
+            vars: result.vars,
+            modelOutput: result.response?.output,
+            modelError: result.response?.error,
+            modelTokens: result.response?.tokenUsage,
+            success: result.success,
+            score: result.score,
+          });
+        }
+        await wandb.finish();
+      }
+
       telemetry.maybeShowNotice();
 
       const border = '='.repeat(process.stdout.columns - 10);

src/types.ts

@@ -21,6 +21,9 @@ export interface CommandLineOptions {
   share?: boolean;
   progressBar?: boolean;
 
+  wandbProject?: string;
+  wandbName?: string;
+
   generateSuggestions?: boolean;
   promptPrefix?: string;
   promptSuffix?: string;