Repository: promptfoo/promptfoo
Pull request #50 (2023-06-28 15:39:07)

Add support for var objects that can be manipulated in nunjucks templates
For example:

template - 
```nunjucks
{% for message in previous_messages %}
{
  "role": "{{message|first}}",
  "content": "{{ message[message|first] }}"
}{% if not loop.last %},{% endif %}
{% endfor %}
```

config -
```yaml
tests:
  - vars:
    previous_messages:
      - user: hello world
      - assistant: how are you?
      - user: great, thanks
```


Patch:
diff --git a/src/evaluator.ts b/src/evaluator.ts
index d877ce9ea..a17f19c76 100644
--- a/src/evaluator.ts
+++ b/src/evaluator.ts
@@ -38,18 +38,24 @@ interface RunEvalOptions {
 const DEFAULT_MAX_CONCURRENCY = 4;
 
 function generateVarCombinations(
-  vars: Record<string, string | string[]>,
-): Record<string, string>[] {
+  vars: Record<string, string | string[] | any>,
+): Record<string, string | any[]>[] {
   const keys = Object.keys(vars);
-  const combinations: Record<string, string>[] = [{}];
+  const combinations: Record<string, string | any[]>[] = [{}];
 
   for (const key of keys) {
-    const values = Array.isArray(vars[key]) ? vars[key] : [vars[key]];
-    const newCombinations: Record<string, string>[] = [];
+    let values: any[] = Array.isArray(vars[key]) ? vars[key] : [vars[key]];
+
+    // Check if it's an array but not a string array
+    if (Array.isArray(vars[key]) && typeof vars[key][0] !== 'string') {
+      values = [vars[key]];
+    }
+
+    const newCombinations: Record<string, any>[] = [];
 
     for (const combination of combinations) {
       for (const value of values) {
-        newCombinations.push({ ...combination, [key]: value as string });
+        newCombinations.push({ ...combination, [key]: value });
       }
     }
 
@@ -229,10 +235,10 @@ class Evaluator {
     });
 
     const varNames: Set<string> = new Set();
-    const varsWithSpecialColsRemoved: Record<string, string | string[]>[] = [];
+    const varsWithSpecialColsRemoved: Record<string, string | string[] | object>[] = [];
     for (const testCase of tests) {
       if (testCase.vars) {
-        const varWithSpecialColsRemoved: Record<string, string | string[]> = {};
+        const varWithSpecialColsRemoved: Record<string, string | string[] | object> = {};
         for (const varName of Object.keys(testCase.vars)) {
           varNames.add(varName);
           varWithSpecialColsRemoved[varName] = testCase.vars[varName];
@@ -354,13 +360,23 @@ class Evaluator {
           resultText = row.response?.output || row.error || '';
         }
 
-        // TODO(ian): Provide full context in table cells, and have the caller
-        // construct the table contents itself.
         const { rowIndex, colIndex } = options;
         if (!table.body[rowIndex]) {
           table.body[rowIndex] = {
             outputs: [],
-            vars: table.head.vars.map((varName) => options.test.vars?.[varName] || '').flat(),
+            vars: table.head.vars
+              .map((varName) => {
+                const varValue = options.test.vars?.[varName] || '';
+                if (typeof varValue === 'string') {
+                  return varValue;
+                }
+                if (Array.isArray(varValue)) {
+                  // Only flatten string arrays
+                  return typeof varValue[0] === 'string' ? varValue : JSON.stringify(varValue);
+                }
+                return JSON.stringify(varValue);
+              })
+              .flat(),
           };
         }
         table.body[rowIndex].outputs[colIndex] = {
diff --git a/src/types.ts b/src/types.ts
index 769ea760a..b6dba8deb 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -84,7 +84,7 @@ export interface Prompt {
 
 export interface EvaluateResult {
   prompt: Prompt;
-  vars: Record<string, string>;
+  vars: Record<string, string | object>;
   response?: ProviderResponse;
   error?: string;
   success: boolean;
@@ -174,7 +174,7 @@ export interface TestCase {
   description?: string;
 
   // Key-value pairs to substitute in the prompt
-  vars?: Record<string, string | string[]>;
+  vars?: Record<string, string | string[] | object>;
 
   // Optional list of automatic checks to run on the LLM output
   assert?: Assertion[];
@@ -185,7 +185,7 @@ export interface TestCase {
 
 // Same as a TestCase, except the `vars` object has been flattened into its final form.
 export interface AtomicTestCase extends TestCase {
-  vars?: Record<string, string>;
+  vars?: Record<string, string | object>;
 }
 
 // The test suite defines the "knobs" that we are tuning in prompt engineering: providers and prompts
diff --git a/src/web/client/src/ResultsTable.tsx b/src/web/client/src/ResultsTable.tsx
index bbb5eb6bf..edef5d242 100644
--- a/src/web/client/src/ResultsTable.tsx
+++ b/src/web/client/src/ResultsTable.tsx
@@ -54,7 +54,7 @@ interface TruncatedTextProps {
 
 function TruncatedText({ text: rawText, maxLength }: TruncatedTextProps) {
   const [isTruncated, setIsTruncated] = React.useState<boolean>(true);
-  const text = String(rawText);
+  const text = typeof rawText === 'string' ? rawText : JSON.stringify(rawText);
 
   const toggleTruncate = () => {
     setIsTruncated(!isTruncated);
@@ -94,14 +94,14 @@ interface PromptOutputProps {
   onRating: (rowIndex: number, promptIndex: number, isPass: boolean) => void;
 }
 
-function PromptOutput({
+function EvalOutputCell({
   output,
   maxTextLength,
   rowIndex,
   promptIndex,
   onRating,
 }: PromptOutputProps) {
-  let text = String(output.text);
+  let text = typeof output.text === 'string' ? output.text : JSON.stringify(output.text);
   let chunks: string[] = [];
   if (!output.pass && text.includes('---')) {
     // TODO(ian): Plumb through failure message instead of parsing it out.
@@ -199,21 +199,26 @@ export default function ResultsTable({
       id: 'vars',
       header: () => <span>Variables</span>,
       columns: head.vars.map((varName, idx) =>
-        columnHelper.accessor((row: EvalRow) => row.vars[idx], {
-          id: `Variable ${idx + 1}`,
-          header: () => (
-            <TableHeader
-              smallText={`Variable ${idx + 1}`}
-              text={varName}
-              maxLength={maxTextLength}
-            />
-          ),
-          cell: (info: CellContext<EvalRow, string>) => (
-            <TruncatedText text={info.getValue()} maxLength={maxTextLength} />
-          ),
-          // Minimize the size of Variable columns.
-          size: 50,
-        }),
+        columnHelper.accessor(
+          (row: EvalRow) => {
+            return row.vars[idx];
+          },
+          {
+            id: `Variable ${idx + 1}`,
+            header: () => (
+              <TableHeader
+                smallText={`Variable ${idx + 1}`}
+                text={varName}
+                maxLength={maxTextLength}
+              />
+            ),
+            cell: (info: CellContext<EvalRow, string>) => (
+              <TruncatedText text={info.getValue()} maxLength={maxTextLength} />
+            ),
+            // Minimize the size of Variable columns.
+            size: 50,
+          },
+        ),
       ),
     }),
     columnHelper.group({
@@ -258,7 +263,7 @@ export default function ResultsTable({
             );
           },
           cell: (info: CellContext<EvalRow, string>) => (
-            <PromptOutput
+            <EvalOutputCell
               output={info.getValue() as unknown as EvalRowOutput}
               maxTextLength={maxTextLength}
               rowIndex={info.row.index}
diff --git a/src/web/client/src/types.ts b/src/web/client/src/types.ts
index aa8ecf6ef..61ef2b19e 100644
--- a/src/web/client/src/types.ts
+++ b/src/web/client/src/types.ts
@@ -6,7 +6,7 @@ export type EvalHead = {
 export type EvalRowOutput = {
   pass: boolean;
   score: number;
-  text: string;
+  text: string | object;
 };
 
 export type EvalRow = {


Test patch:
diff --git a/test/evaluator.test.ts b/test/evaluator.test.ts
index 8a60cc8da..2d705fab8 100644
--- a/test/evaluator.test.ts
+++ b/test/evaluator.test.ts
@@ -61,6 +61,28 @@ describe('evaluator', () => {
     expect(summary.results[0].response?.output).toBe('Test output');
   });
 
+  test('evaluate with vars as object', async () => {
+    const testSuite: TestSuite = {
+      providers: [mockApiProvider],
+      prompts: [toPrompt('Test prompt {{ var1.prop1 }} {{ var2 }}')],
+      tests: [
+        {
+          vars: { var1: { prop1: 'value1' }, var2: 'value2' },
+        },
+      ],
+    };
+
+    const summary = await evaluate(testSuite, {});
+
+    expect(mockApiProvider.callApi).toHaveBeenCalledTimes(1);
+    expect(summary.stats.successes).toBe(1);
+    expect(summary.stats.failures).toBe(0);
+    expect(summary.stats.tokenUsage).toEqual({ total: 10, prompt: 5, completion: 5, cached: 0 });
+    expect(summary.results[0].prompt.raw).toBe('Test prompt value1 value2');
+    expect(summary.results[0].prompt.display).toBe('Test prompt {{ var1.prop1 }} {{ var2 }}');
+    expect(summary.results[0].response?.output).toBe('Test output');
+  });
+
   test('evaluate with named prompt', async () => {
     const testSuite: TestSuite = {
       providers: [mockApiProvider],