Repository: promptfoo/promptfoo
Pull request #93 (2023-07-27 05:49:31)

Add support for native function ApiProviders and Assertions
This change is for people using promptfoo as a node package.  It enables you to:
- Specify a provider that is a function, so you can call your Javascript code directly
- Specify a `value` for the `javascript` assert type that is a function, so you can include custom logic more easily instead of stuffing it into a string.

For example:

```js
    providers: [
      'openai:gpt-3.5-turbo',
      (prompt) => {
        // Call LLM here...
        return {
          output: '<LLM output>',
        };
      },
    ],
    tests: [
      {
        vars: {
          body: "I'm hungry",
        },
        assert: [
          {
            type: 'javascript',
            value: (output) => {
              const pass = output.includes("J'ai faim");
              return {
                pass,
                score: pass ? 1.0 : 0.0,
                reason: pass ? 'Output contained substring' : 'Output did not contain substring',
              };
            },
          },
        ],
      },
    ],
```

#90 


Patch:
diff --git a/examples/node-package/index.js b/examples/node-package/index.js
index 56266af11..60514bcc2 100644
--- a/examples/node-package/index.js
+++ b/examples/node-package/index.js
@@ -3,7 +3,15 @@ import promptfoo from '../../dist/src/index.js';
 (async () => {
   const results = await promptfoo.evaluate({
     prompts: ['Rephrase this in French: {{body}}', 'Rephrase this like a pirate: {{body}}'],
-    providers: ['openai:gpt-3.5-turbo'],
+    providers: [
+      'openai:gpt-3.5-turbo',
+      (prompt) => {
+        // Call LLM here...
+        return {
+          output: '<LLM output>',
+        };
+      },
+    ],
     tests: [
       {
         vars: {
@@ -14,6 +22,19 @@ import promptfoo from '../../dist/src/index.js';
         vars: {
           body: "I'm hungry",
         },
+        assert: [
+          {
+            type: 'javascript',
+            value: (output) => {
+              const pass = output.includes("J'ai faim");
+              return {
+                pass,
+                score: pass ? 1.0 : 0.0,
+                reason: pass ? 'Output contained substring' : 'Output did not contain substring',
+              };
+            },
+          },
+        ],
       },
     ],
   });
diff --git a/src/assertions.ts b/src/assertions.ts
index 9cebfa042..261946e88 100644
--- a/src/assertions.ts
+++ b/src/assertions.ts
@@ -234,12 +234,16 @@ export async function runAssertion(
     };
   }
 
+  const context = {
+    vars: test.vars || {},
+  };
+
   if (baseType === 'javascript') {
     try {
+      if (typeof assertion.value === 'function') {
+        return assertion.value(output, test, assertion);
+      }
       const customFunction = new Function('output', 'context', `return ${assertion.value}`);
-      const context = {
-        vars: test.vars || {},
-      };
       const result = customFunction(output, context) as any;
       if (typeof result === 'boolean') {
         pass = result !== inverse;
@@ -268,6 +272,44 @@ ${assertion.value}`,
     };
   }
 
+  if (baseType === 'python') {
+    try {
+      const { execSync } = require('child_process');
+      const escapedOutput = output.replace(/'/g, "\\'").replace(/"/g, '\\"');;
+      const escapedContext = JSON.stringify(context).replace(/'/g, "\\'").replace(/"/g, '\\"');
+      const result = execSync(`python -c "import json; import math; import os; import sys; import re; import datetime; import random; import collections; output='${escapedOutput}'; context='${escapedContext}'; print(json.dumps(${assertion.value}))"`).toString().trim();
+      if (result === 'true') {
+        pass = true;
+        score = 1.0;
+      } else if (result === 'false') {
+        pass = false;
+        score = 0.0;
+      } else if (result.startsWith('{')) {
+        return JSON.parse(result);
+      } else {
+        pass = true;
+        score = parseFloat(result);
+        if (isNaN(score)) {
+          throw new Error('Python code must return a boolean, number, or {pass, score, reason} object');
+        }
+      }
+    } catch (err) {
+      return {
+        pass: false,
+        score: 0,
+        reason: `Python code execution failed: ${(err as Error).message}`,
+      };
+    }
+    return {
+      pass,
+      score,
+      reason: pass
+        ? 'Assertion passed'
+        : `Python code returned ${pass ? 'true' : 'false'}
+${assertion.value}`,
+    };
+  }
+
   if (baseType === 'similar') {
     invariant(assertion.value, 'Similarity assertion must have a string value');
     invariant(
@@ -339,7 +381,10 @@ ${assertion.value}`,
   }
 
   if (baseType === 'rouge-n') {
-    invariant(assertion.value, '"rouge" assertion type must a value (string or string array)');
+    invariant(
+      typeof assertion.value === 'string' || Array.isArray(assertion.value),
+      '"rouge" assertion type must be a value (string or string array)',
+    );
     return handleRougeScore(baseType, assertion, assertion.value, output, inverse);
   }
 
diff --git a/src/providers.ts b/src/providers.ts
index a99d28b6f..4caf62248 100644
--- a/src/providers.ts
+++ b/src/providers.ts
@@ -1,7 +1,5 @@
 import path from 'path';
 
-import { ApiProvider, ProviderConfig, ProviderId, RawProviderConfig } from './types';
-
 import { OpenAiCompletionProvider, OpenAiChatCompletionProvider } from './providers/openai';
 import { AnthropicCompletionProvider } from './providers/anthropic';
 import { ReplicateProvider } from './providers/replicate';
@@ -12,17 +10,37 @@ import {
   AzureOpenAiCompletionProvider,
 } from './providers/azureopenai';
 
+import type {
+  ApiProvider,
+  ProviderConfig,
+  ProviderFunction,
+  ProviderId,
+  RawProviderConfig,
+} from './types';
+
 export async function loadApiProviders(
-  providerPaths: ProviderId | ProviderId[] | RawProviderConfig[],
+  providerPaths: ProviderId | ProviderId[] | RawProviderConfig[] | ProviderFunction,
   basePath?: string,
 ): Promise<ApiProvider[]> {
   if (typeof providerPaths === 'string') {
     return [await loadApiProvider(providerPaths, undefined, basePath)];
+  } else if (typeof providerPaths === 'function') {
+    return [
+      {
+        id: () => 'custom-function',
+        callApi: providerPaths,
+      },
+    ];
   } else if (Array.isArray(providerPaths)) {
     return Promise.all(
-      providerPaths.map((provider) => {
+      providerPaths.map((provider, idx) => {
         if (typeof provider === 'string') {
           return loadApiProvider(provider, undefined, basePath);
+        } else if (typeof provider === 'function') {
+          return {
+            id: () => `custom-function-${idx}`,
+            callApi: provider,
+          };
         } else {
           const id = Object.keys(provider)[0];
           const context = { ...provider[id], id };
diff --git a/src/types.ts b/src/types.ts
index 23caa68d8..f9c69a48e 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -151,6 +151,7 @@ type BaseAssertionTypes =
   | 'is-json'
   | 'contains-json'
   | 'javascript'
+  | 'python'
   | 'similar'
   | 'llm-rubric'
   | 'webhook'
@@ -168,7 +169,10 @@ export interface Assertion {
   type: AssertionType;
 
   // The expected value, if applicable
-  value?: string | string[];
+  value?:
+    | string
+    | string[]
+    | ((output: string, testCase: AtomicTestCase, assertion: Assertion) => Promise<GradingResult>);
 
   // The threshold value, only applicable for similarity (cosine distance)
   threshold?: number;
@@ -227,6 +231,8 @@ export interface TestSuite {
 
 export type ProviderId = string;
 
+export type ProviderFunction = (prompt: string) => Promise<ProviderResponse>;
+
 export type RawProviderConfig = Record<ProviderId, Omit<ProviderConfig, 'id'>>;
 
 // TestSuiteConfig = Test Suite, but before everything is parsed and resolved.  Providers are just strings, prompts are filepaths, tests can be filepath or inline.
@@ -235,7 +241,7 @@ export interface TestSuiteConfig {
   description?: string;
 
   // One or more LLM APIs to use, for example: openai:gpt-3.5-turbo, openai:gpt-4, localai:chat:vicuna
-  providers: ProviderId | ProviderId[] | RawProviderConfig[];
+  providers: ProviderId | ProviderId[] | RawProviderConfig[] | ProviderFunction;
 
   // One or more prompt files to load
   prompts: string | string[];
diff --git a/src/util.ts b/src/util.ts
index bc5b369f3..640653673 100644
--- a/src/util.ts
+++ b/src/util.ts
@@ -4,6 +4,7 @@ import * as os from 'os';
 
 import $RefParser from '@apidevtools/json-schema-ref-parser';
 import fetch from 'node-fetch';
+import invariant from 'tiny-invariant';
 import yaml from 'js-yaml';
 import nunjucks from 'nunjucks';
 import { globSync } from 'glob';
@@ -44,6 +45,15 @@ export function readProviderPromptMap(
     allPrompts.push(prompt.display);
   }
 
+  invariant(
+    typeof config.providers !== 'string',
+    'In order to use a provider-prompt map, config.providers should be an array of objects, not a string',
+  );
+  invariant(
+    typeof config.providers !== 'function',
+    'In order to use a provider-prompt map, config.providers should be an array of objects, not a function',
+  );
+
   for (const provider of config.providers) {
     if (typeof provider === 'object') {
       const rawProvider = provider as RawProviderConfig;


Test patch:
diff --git a/test/assertions.test.ts b/test/assertions.test.ts
index 3bb15f94c..6c0cb99e7 100644
--- a/test/assertions.test.ts
+++ b/test/assertions.test.ts
@@ -58,16 +58,26 @@ describe('runAssertion', () => {
     type: 'contains-json',
   };
 
-  const functionAssertion: Assertion = {
+  const javascriptStringAssertion: Assertion = {
     type: 'javascript',
     value: 'output === "Expected output"',
   };
 
-  const functionAssertionWithNumber: Assertion = {
+  const javascriptStringAssertionWithNumber: Assertion = {
     type: 'javascript',
     value: 'output.length * 10',
   };
 
+  const javascriptFunctionAssertion: Assertion = {
+    type: 'javascript',
+    value: async (output: string) => ({ pass: true, score: 0.5, reason: 'Assertion passed' }),
+  };
+
+  const javascriptFunctionFailAssertion: Assertion = {
+    type: 'javascript',
+    value: async (output: string) => ({ pass: false, score: 0.5, reason: 'Assertion failed' }),
+  };
+
   it('should pass when the equality assertion passes', async () => {
     const output = 'Expected output';
 
@@ -133,11 +143,11 @@ describe('runAssertion', () => {
     expect(result.reason).toContain('Expected output to contain valid JSON');
   });
 
-  it('should pass when the function assertion passes', async () => {
+  it('should pass when the javascript assertion passes', async () => {
     const output = 'Expected output';
 
     const result: GradingResult = await runAssertion(
-      functionAssertion,
+      javascriptStringAssertion,
       {} as AtomicTestCase,
       output,
     );
@@ -145,11 +155,11 @@ describe('runAssertion', () => {
     expect(result.reason).toBe('Assertion passed');
   });
 
-  it('should pass a score through when the function returns a number', async () => {
+  it('should pass a score through when the javascript returns a number', async () => {
     const output = 'Expected output';
 
     const result: GradingResult = await runAssertion(
-      functionAssertionWithNumber,
+      javascriptStringAssertionWithNumber,
       {} as AtomicTestCase,
       output,
     );
@@ -158,11 +168,11 @@ describe('runAssertion', () => {
     expect(result.reason).toBe('Assertion passed');
   });
 
-  it('should fail when the function assertion fails', async () => {
+  it('should fail when the javascript assertion fails', async () => {
     const output = 'Different output';
 
     const result: GradingResult = await runAssertion(
-      functionAssertion,
+      javascriptStringAssertion,
       {} as AtomicTestCase,
       output,
     );
@@ -173,12 +183,12 @@ describe('runAssertion', () => {
   it('should pass when the function assertion passes - with vars', async () => {
     const output = 'Expected output';
 
-    const functionAssertionWithVars: Assertion = {
+    const javascriptStringAssertionWithVars: Assertion = {
       type: 'javascript',
       value: 'output === "Expected output" && context.vars.foo === "bar"',
     };
     const result: GradingResult = await runAssertion(
-      functionAssertionWithVars,
+      javascriptStringAssertionWithVars,
       { vars: { foo: 'bar' } } as AtomicTestCase,
       output,
     );
@@ -186,15 +196,15 @@ describe('runAssertion', () => {
     expect(result.reason).toBe('Assertion passed');
   });
 
-  it('should fail when the function does not match vars', async () => {
+  it('should fail when the javascript does not match vars', async () => {
     const output = 'Expected output';
 
-    const functionAssertionWithVars: Assertion = {
+    const javascriptStringAssertionWithVars: Assertion = {
       type: 'javascript',
       value: 'output === "Expected output" && context.vars.foo === "something else"',
     };
     const result: GradingResult = await runAssertion(
-      functionAssertionWithVars,
+      javascriptStringAssertionWithVars,
       { vars: { foo: 'bar' } } as AtomicTestCase,
       output,
     );
@@ -204,6 +214,32 @@ describe('runAssertion', () => {
     );
   });
 
+  it('should pass when the function returns pass', async () => {
+    const output = 'Expected output';
+
+    const result: GradingResult = await runAssertion(
+      javascriptFunctionAssertion,
+      {} as AtomicTestCase,
+      output,
+    );
+    expect(result.pass).toBeTruthy();
+    expect(result.score).toBe(0.5);
+    expect(result.reason).toBe('Assertion passed');
+  });
+
+  it('should fail when the function returns fail', async () => {
+    const output = 'Expected output';
+
+    const result: GradingResult = await runAssertion(
+      javascriptFunctionFailAssertion,
+      {} as AtomicTestCase,
+      output,
+    );
+    expect(result.pass).toBeFalsy();
+    expect(result.score).toBe(0.5);
+    expect(result.reason).toBe('Assertion failed');
+  });
+
   const notContainsAssertion: Assertion = {
     type: 'not-contains',
     value: 'Unexpected output',
diff --git a/test/providers.test.ts b/test/providers.test.ts
index febeff713..1dc7c018c 100644
--- a/test/providers.test.ts
+++ b/test/providers.test.ts
@@ -5,7 +5,7 @@ import { AnthropicCompletionProvider } from '../src/providers/anthropic';
 
 import { disableCache, enableCache } from '../src/cache.js';
 import { loadApiProvider, loadApiProviders } from '../src/providers.js';
-import type { RawProviderConfig } from '../src/types';
+import type { RawProviderConfig, ProviderFunction } from '../src/types';
 import {
   AzureOpenAiChatCompletionProvider,
   AzureOpenAiCompletionProvider,
@@ -201,6 +201,21 @@ describe('providers', () => {
     expect(provider).toBeInstanceOf(OpenAiChatCompletionProvider);
   });
 
+  test('loadApiProviders with ProviderFunction', async () => {
+    const providerFunction: ProviderFunction = async (prompt: string) => {
+      return {
+        output: `Output for ${prompt}`,
+        tokenUsage: { total: 10, prompt: 5, completion: 5 },
+      };
+    };
+    const providers = await loadApiProviders(providerFunction);
+    expect(providers).toHaveLength(1);
+    expect(providers[0].id()).toBe('custom-function');
+    const response = await providers[0].callApi('Test prompt');
+    expect(response.output).toBe('Output for Test prompt');
+    expect(response.tokenUsage).toEqual({ total: 10, prompt: 5, completion: 5 });
+  });
+
   test('loadApiProviders with RawProviderConfig[]', async () => {
     const rawProviderConfigs: RawProviderConfig[] = [
       {
diff --git a/test/util.test.ts b/test/util.test.ts
index 77b074043..78f47fa2e 100644
--- a/test/util.test.ts
+++ b/test/util.test.ts
@@ -496,6 +496,6 @@ describe('readTests', () => {
     const result = await readTests(testsPaths);
 
     expect(fs.readFileSync).toHaveBeenCalledTimes(2);
-    expect(result).toEqual([Object.assign({}, test1Content[0], {vars: vars1Content})]);
+    expect(result).toEqual([Object.assign({}, test1Content[0], { vars: vars1Content })]);
   });
 });