promptfoo 0.8.3 → 0.10.0

package/src/util.ts

@@ -10,19 +10,13 @@ import { parse as parsePath } from 'path';
 import { parse as parseCsv } from 'csv-parse/sync';
 import { stringify } from 'csv-stringify/sync';
 
-import logger from './logger.js';
-import { getDirectory } from './esm.js';
+import logger from './logger';
+import { getDirectory } from './esm';
 
 import type { RequestInfo, RequestInit, Response } from 'node-fetch';
 
-import type {
-  Assertion,
-  CsvRow,
-  EvaluateSummary,
-  UnifiedConfig,
-  TestCase,
-} from './types.js';
-import { assertionFromString } from './assertions.js';
+import type { Assertion, CsvRow, EvaluateSummary, UnifiedConfig, TestCase, Prompt } from './types';
+import { assertionFromString } from './assertions';
 
 const PROMPT_DELIMITER = '---';
 
@@ -50,17 +44,36 @@ export function readConfig(configPath: string): UnifiedConfig {
     case '.js':
       return require(configPath) as UnifiedConfig;
     case '.yaml':
+    case '.yml':
       return yaml.load(fs.readFileSync(configPath, 'utf-8')) as UnifiedConfig;
     default:
       throw new Error(`Unsupported configuration file format: ${ext}`);
   }
 }
 
-export function readPrompts(promptPathsOrGlobs: string | string[]): string[] {
-  promptPathsOrGlobs =
-    typeof promptPathsOrGlobs === 'string' ? [promptPathsOrGlobs] : promptPathsOrGlobs;
-  const promptPaths = promptPathsOrGlobs.flatMap((pathOrGlob) => globSync(pathOrGlob));
-  let promptContents: string[] = [];
+enum PromptInputType {
+  STRING = 1,
+  ARRAY = 2,
+  NAMED = 3,
+}
+
+export function readPrompts(
+  promptPathOrGlobs: string | string[] | Record<string, string>,
+): Prompt[] {
+  let promptPaths: string[] = [];
+  let promptContents: Prompt[] = [];
+
+  let inputType: PromptInputType | undefined;
+  if (typeof promptPathOrGlobs === 'string') {
+    promptPaths = [promptPathOrGlobs];
+    inputType = PromptInputType.STRING;
+  } else if (Array.isArray(promptPathOrGlobs)) {
+    promptPaths = promptPathOrGlobs.flatMap((pathOrGlob) => globSync(pathOrGlob));
+    inputType = PromptInputType.ARRAY;
+  } else if (typeof promptPathOrGlobs === 'object') {
+    promptPaths = Object.keys(promptPathOrGlobs);
+    inputType = PromptInputType.NAMED;
+  }
 
   for (const promptPath of promptPaths) {
     const stat = fs.statSync(promptPath);
@@ -69,18 +82,27 @@ export function readPrompts(promptPathsOrGlobs: string | string[]): string[] {
       const fileContents = filesInDirectory.map((fileName) =>
         fs.readFileSync(path.join(promptPath, fileName), 'utf-8'),
       );
-      promptContents.push(...fileContents);
+      promptContents.push(...fileContents.map((content) => ({ raw: content, display: content })));
     } else {
       const fileContent = fs.readFileSync(promptPath, 'utf-8');
-      promptContents.push(fileContent);
+      let display;
+      if (inputType === PromptInputType.NAMED) {
+        display = (promptPathOrGlobs as Record<string, string>)[promptPath];
+      } else {
+        display = fileContent.length > 200 ? promptPath : fileContent;
+      }
+      promptContents.push({ raw: fileContent, display });
     }
   }
 
-  if (promptContents.length === 1) {
-    promptContents = promptContents[0].split(PROMPT_DELIMITER).map((p) => p.trim());
+  if (promptContents.length === 1 && inputType !== PromptInputType.NAMED) {
+    const content = promptContents[0].raw;
+    promptContents = content
+      .split(PROMPT_DELIMITER)
+      .map((p) => ({ raw: p.trim(), display: p.trim() }));
   }
   if (promptContents.length === 0) {
-    throw new Error(`There are no prompts in ${promptPathsOrGlobs.join(', ')}`);
+    throw new Error(`There are no prompts in ${JSON.stringify(promptPathOrGlobs)}`);
   }
   return promptContents;
 }