Repository: promptfoo/promptfoo
Pull request #84 (2023-07-22 03:15:07)

Add support for postprocessing in test cases
The `TestCase.options.postprocess` field is a Javascript snippet that modifies the LLM output.  Postprocessing occurs before any assertions are run.

For example:

```yaml
# ...
tests:
  - vars:
      language: French
      body: Hello world
    options:
      postprocess: output.toUpperCase()
    # ...
```

Or multiline:

```yaml
# ...
tests:
  - vars:
      language: French
      body: Hello world
    options:
      postprocess: |
        const words = output.split(' ').filter(x => !!x);
        return JSON.stringify(words);
    # ...
```


Patch:
diff --git a/src/evaluator.ts b/src/evaluator.ts
index 2a3dd92ff..26f1b3fd2 100644
--- a/src/evaluator.ts
+++ b/src/evaluator.ts
@@ -3,6 +3,7 @@ import readline from 'readline';
 import async from 'async';
 import chalk from 'chalk';
 import nunjucks from 'nunjucks';
+import invariant from 'tiny-invariant';
 
 import logger from './logger';
 import telemetry from './telemetry';
@@ -121,7 +122,23 @@ class Evaluator {
       if (response.error) {
         ret.error = response.error;
       } else if (response.output) {
-        const checkResult = await runAssertions(test, response.output);
+        // Create a copy of response so we can potentially mutate it.
+        let processedResponse = { ...response };
+        if (test.options?.postprocess) {
+          const { postprocess } = test.options;
+          const postprocessFn = new Function(
+            'output',
+            'context',
+            postprocess.includes('\n') ? postprocess : `return ${postprocess}`,
+          );
+          processedResponse.output = postprocessFn(processedResponse.output);
+          if (processedResponse.output == null) {
+            throw new Error('Postprocess function did not return a value');
+          }
+        }
+
+        invariant(processedResponse.output != null, 'Response output should not be null');
+        const checkResult = await runAssertions(test, processedResponse.output);
         if (!checkResult.pass) {
           ret.error = checkResult.reason;
         }
@@ -132,6 +149,7 @@ class Evaluator {
           this.stats.tokenUsage.prompt += checkResult.tokensUsed.prompt;
           this.stats.tokenUsage.completion += checkResult.tokensUsed.completion;
         }
+        ret.response = processedResponse;
       } else {
         ret.success = false;
         ret.score = 0;
@@ -272,6 +290,8 @@ class Evaluator {
         testCase.options?.prefix || testSuite.defaultTest?.options?.prefix || '';
       const appendToPrompt =
         testCase.options?.suffix || testSuite.defaultTest?.options?.suffix || '';
+      testCase.options.postprocess =
+        testCase.options.postprocess || testSuite.defaultTest?.options?.postprocess;
 
       // Finalize test case eval
       const varCombinations = generateVarCombinations(testCase.vars || {});
diff --git a/src/main.ts b/src/main.ts
index 8a6027577..6f63f3323 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -323,6 +323,7 @@ async function main() {
           suffix: cmdObj.promptSuffix,
           provider: cmdObj.grader,
           // rubricPrompt:
+          // postprocess
         },
         ...config.defaultTest,
       };
diff --git a/src/types.ts b/src/types.ts
index 53431ffd0..396d55b54 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -73,6 +73,10 @@ export interface PromptConfig {
   suffix?: string;
 }
 
+export interface OutputConfig {
+  postprocess?: string;
+}
+
 export interface EvaluateOptions {
   maxConcurrency?: number;
   showProgressBar?: boolean;
@@ -185,7 +189,7 @@ export interface TestCase {
   assert?: Assertion[];
 
   // Additional configuration settings for the prompt
-  options?: PromptConfig & GradingConfig;
+  options?: PromptConfig & OutputConfig & GradingConfig;
 }
 
 // Same as a TestCase, except the `vars` object has been flattened into its final form.


Test patch:
diff --git a/test/evaluator.test.ts b/test/evaluator.test.ts
index 46317a001..945b38f4b 100644
--- a/test/evaluator.test.ts
+++ b/test/evaluator.test.ts
@@ -345,6 +345,52 @@ describe('evaluator', () => {
     expect(summary.results[0].response?.output).toBe('Test output');
   });
 
+  test('evaluate with postprocess option - default test', async () => {
+    const testSuite: TestSuite = {
+      providers: [mockApiProvider],
+      prompts: [toPrompt('Test prompt')],
+      defaultTest: {
+        options: {
+          postprocess: 'output + " postprocessed"',
+        },
+      },
+    };
+
+    const summary = await evaluate(testSuite, {});
+
+    expect(mockApiProvider.callApi).toHaveBeenCalledTimes(1);
+    expect(summary.stats.successes).toBe(1);
+    expect(summary.stats.failures).toBe(0);
+    expect(summary.results[0].response?.output).toBe('Test output postprocessed');
+  });
+
+  test('evaluate with postprocess option - single test', async () => {
+    const testSuite: TestSuite = {
+      providers: [mockApiProvider],
+      prompts: [toPrompt('Test prompt')],
+      tests: [
+        {
+          assert: [
+            {
+              type: 'equals',
+              value: 'Test output postprocessed',
+            },
+          ],
+          options: {
+            postprocess: 'output + " postprocessed"',
+          },
+        },
+      ],
+    };
+
+    const summary = await evaluate(testSuite, {});
+
+    expect(mockApiProvider.callApi).toHaveBeenCalledTimes(1);
+    expect(summary.stats.successes).toBe(1);
+    expect(summary.stats.failures).toBe(0);
+    expect(summary.results[0].response?.output).toBe('Test output postprocessed');
+  });
+
   test('evaluate with providerPromptMap', async () => {
     const testSuite: TestSuite = {
       providers: [mockApiProvider],