promptfoo 0.1.0

package/src/evaluator.ts

@@ -0,0 +1,203 @@
+import async from 'async';
+import nunjucks from 'nunjucks';
+
+import type { SingleBar } from 'cli-progress';
+
+import { EvaluateOptions, EvaluateSummary, EvaluateResult, ApiProvider, Prompt } from './types.js';
+
+interface RunEvalOptions {
+  provider: ApiProvider;
+  prompt: string;
+  vars?: Record<string, string>;
+  includeProviderId?: boolean;
+}
+
+const DEFAULT_MAX_CONCURRENCY = 3;
+
+function checkExpectedValue(expected: string, output: string): boolean {
+  if (expected.startsWith('eval:')) {
+    const evalBody = expected.slice(5);
+    const evalFunction = new Function('output', `return ${evalBody}`);
+    return evalFunction(output);
+  } else if (expected.startsWith('grade:')) {
+    // NYI
+    return false;
+  } else {
+    return expected === output;
+  }
+}
+
+async function runEval({
+  provider,
+  prompt,
+  vars,
+  includeProviderId,
+}: RunEvalOptions): Promise<EvaluateResult> {
+  vars = vars || {};
+  const renderedPrompt = nunjucks.renderString(prompt, vars);
+
+  // Note that we're using original prompt, not renderedPrompt
+  const promptDisplay = includeProviderId ? `[${provider.id()}] ${prompt}` : prompt;
+
+  const setup = {
+    prompt: {
+      raw: renderedPrompt,
+      display: promptDisplay,
+    },
+    vars,
+  };
+
+  try {
+    const response = await provider.callApi(renderedPrompt);
+    const success = vars.__expected ? checkExpectedValue(vars.__expected, response.output) : true;
+    const ret: EvaluateResult = {
+      ...setup,
+      response,
+      success,
+    };
+    if (!success) {
+      ret.error = `Expected ${vars.__expected}, got "${response.output}"`;
+    }
+    return ret;
+  } catch (err) {
+    return {
+      ...setup,
+      error: String(err),
+      success: false,
+    };
+  }
+}
+
+export async function evaluate(options: EvaluateOptions): Promise<EvaluateSummary> {
+  const prompts: Prompt[] = [];
+  const results: EvaluateResult[] = [];
+
+  for (const promptContent of options.prompts) {
+    for (const provider of options.providers) {
+      prompts.push({
+        raw: promptContent,
+        display:
+          options.providers.length > 1 ? `[${provider.id()}] ${promptContent}` : promptContent,
+      });
+    }
+  }
+
+  const vars = options.vars && options.vars.length > 0 ? options.vars : [{}];
+  const varsWithExpectedKeyRemoved = vars.map((v) => {
+    const ret = { ...v };
+    delete ret.__expected;
+    return ret;
+  });
+  const isTest = vars[0].__expected;
+  const table: string[][] = [
+    isTest
+      ? [
+          'RESULT',
+          [...prompts.map((p) => p.display), ...Object.keys(varsWithExpectedKeyRemoved[0])],
+        ].flat()
+      : [...prompts.map((p) => p.display), ...Object.keys(varsWithExpectedKeyRemoved[0])],
+  ];
+
+  const stats = {
+    successes: 0,
+    failures: 0,
+    tokenUsage: {
+      total: 0,
+      prompt: 0,
+      completion: 0,
+    },
+  };
+
+  let progressbar: SingleBar | undefined;
+  if (options.showProgressBar) {
+    const totalNumRuns =
+      options.prompts.length * options.providers.length * (options.vars?.length || 1);
+    const cliProgress = await import('cli-progress');
+    progressbar = new cliProgress.SingleBar(
+      {
+        format:
+          'Eval: [{bar}] {percentage}% | ETA: {eta}s | {value}/{total} | {provider} "{prompt}" {vars}',
+      },
+      cliProgress.Presets.shades_classic,
+    );
+    progressbar.start(totalNumRuns, 0, {
+      provider: '',
+      prompt: '',
+      vars: '',
+    });
+  }
+
+  const runEvalOptions: RunEvalOptions[] = [];
+  for (const row of vars) {
+    for (const promptContent of options.prompts) {
+      for (const provider of options.providers) {
+        runEvalOptions.push({
+          provider,
+          prompt: promptContent,
+          vars: row,
+          includeProviderId: options.providers.length > 1,
+        });
+      }
+    }
+  }
+
+  const combinedOutputs: string[][] = new Array(vars.length).fill(null).map(() => []);
+  await async.forEachOfLimit(
+    runEvalOptions,
+    options.maxConcurrency || DEFAULT_MAX_CONCURRENCY,
+    async (options: RunEvalOptions, index: number | string) => {
+      const row = await runEval(options);
+      results.push(row);
+      if (row.error) {
+        stats.failures++;
+      } else {
+        if (row.success) {
+          stats.successes++;
+        } else {
+          stats.failures++;
+        }
+        stats.tokenUsage.total += row.response?.tokenUsage?.total || 0;
+        stats.tokenUsage.prompt += row.response?.tokenUsage?.prompt || 0;
+        stats.tokenUsage.completion += row.response?.tokenUsage?.completion || 0;
+      }
+
+      if (progressbar) {
+        progressbar.increment({
+          provider: options.provider.id(),
+          prompt: options.prompt.slice(0, 10),
+          vars: Object.entries(options.vars || {})
+            .map(([k, v]) => `${k}=${v}`)
+            .join(' ')
+            .slice(0, 10),
+        });
+      }
+
+      // Bookkeeping for table
+      if (typeof index !== 'number') {
+        throw new Error('Expected index to be a number');
+      }
+      const combinedOutputIndex = Math.floor(index / prompts.length);
+      combinedOutputs[combinedOutputIndex].push(row.response?.output || '');
+    },
+  );
+
+  if (progressbar) {
+    progressbar.stop();
+  }
+
+  if (isTest) {
+    table.push(
+      ...combinedOutputs.map((output, index) => [
+        results[index].success ? 'PASS' : `FAIL: ${results[index].error}`,
+        ...output,
+        ...Object.values(varsWithExpectedKeyRemoved[index]),
+      ]),
+    );
+  } else {
+    table.push(
+      ...combinedOutputs.map((output, index) => [...output, ...Object.values(vars[index])]),
+    );
+  }
+
+  return { results, stats, table };
+}

package/src/index.ts

@@ -0,0 +1,35 @@
+import { evaluate as doEvaluate } from './evaluator.js';
+import { loadApiProvider } from './providers.js';
+
+import type { ApiProvider, EvaluateOptions, EvaluateSummary } from './types.js';
+
+async function evaluate(
+  providers: (string | ApiProvider)[] | (string | ApiProvider),
+  options: Omit<EvaluateOptions, 'providers'>,
+): Promise<EvaluateSummary> {
+  let apiProviders: ApiProvider[] = [];
+  const addProvider = async (provider: ApiProvider | string) => {
+    if (typeof provider === 'string') {
+      apiProviders.push(await loadApiProvider(provider));
+    } else {
+      apiProviders.push(provider);
+    }
+  };
+
+  if (Array.isArray(providers)) {
+    for (const provider of providers) {
+      await addProvider(provider);
+    }
+  } else {
+    await addProvider(providers);
+  }
+
+  return doEvaluate({
+    ...options,
+    providers: apiProviders,
+  });
+}
+
+export default {
+  evaluate,
+};

package/src/logger.ts

@@ -0,0 +1,38 @@
+import chalk from 'chalk';
+import winston from 'winston';
+
+const logLevels = {
+  error: 0,
+  warn: 1,
+  info: 2,
+  debug: 3,
+};
+
+const customFormatter = winston.format.printf(({ level, message, ...args }) => {
+  if (level === 'error') {
+    return chalk.red(message);
+  } else if (level === 'warn') {
+    return chalk.yellow(message);
+  } else if (level === 'info') {
+    return message;
+  } else if (level === 'debug') {
+    return chalk.cyan(message);
+  }
+  throw new Error(`Invalid log level: ${level}`);
+});
+
+const logger = winston.createLogger({
+  levels: logLevels,
+  format: winston.format.combine(winston.format.simple(), customFormatter),
+  transports: [new winston.transports.Console()],
+});
+
+export function setLogLevel(level: keyof typeof logLevels) {
+  if (logLevels.hasOwnProperty(level)) {
+    logger.transports[0].level = level;
+  } else {
+    throw new Error(`Invalid log level: ${level}`);
+  }
+}
+
+export default logger;

package/src/main.ts

@@ -0,0 +1,108 @@
+#!/usr/bin/env node
+import { readFileSync } from 'fs';
+import { parse } from 'path';
+
+import Table from 'cli-table3';
+import chalk from 'chalk';
+import { Command } from 'commander';
+
+import logger, { setLogLevel } from './logger.js';
+import { loadApiProvider } from './providers.js';
+import { evaluate } from './evaluator.js';
+import { readPrompts, readVars, writeOutput } from './util.js';
+
+import type { CommandLineOptions, EvaluateOptions, VarMapping } from './types.js';
+
+const program = new Command();
+
+program
+  .command('eval')
+  .description('Evaluate prompts')
+  .requiredOption('-p, --prompt <paths...>', 'Paths to prompt files (.txt)')
+  .requiredOption(
+    '-r, --provider <name or path...>',
+    'One of: openai:chat, openai:completion, openai:<model name>, or path to custom API caller module',
+  )
+  .option('-o, --output <path>', 'Path to output file (csv, json, yaml, html)')
+  .option('-v, --vars <path>', 'Path to file with prompt variables (csv, json, yaml)')
+  .option('-c, --config <path>', 'Path to configuration file')
+  .option('-j, --max-concurrency <number>', 'Maximum number of concurrent API calls')
+  .option('--verbose', 'Show debug logs')
+  .action(async (cmdObj: CommandLineOptions & Command) => {
+    if (cmdObj.verbose) {
+      setLogLevel('debug');
+    }
+
+    const configPath = cmdObj.config;
+    let config = {};
+    if (configPath) {
+      const ext = parse(configPath).ext;
+      switch (ext) {
+        case '.json':
+          const content = readFileSync(configPath, 'utf-8');
+          config = JSON.parse(content);
+          break;
+        case '.js':
+          config = require(configPath);
+          break;
+        default:
+          throw new Error(`Unsupported configuration file format: ${ext}`);
+      }
+    }
+
+    let vars: VarMapping[] = [];
+    if (cmdObj.vars) {
+      vars = readVars(cmdObj.vars);
+    }
+
+    const providers = cmdObj.provider.map((p) => loadApiProvider(p));
+    const options: EvaluateOptions = {
+      prompts: readPrompts(cmdObj.prompt),
+      vars,
+      providers,
+      showProgressBar: true,
+      maxConcurrency:
+        cmdObj.maxConcurrency && cmdObj.maxConcurrency > 0 ? cmdObj.maxConcurrency : undefined,
+      ...config,
+    };
+
+    const summary = await evaluate(options);
+
+    if (cmdObj.output) {
+      logger.info(chalk.yellow(`Writing output to ${cmdObj.output}`));
+      writeOutput(cmdObj.output, summary);
+    } else {
+      // Output table by default
+      const maxWidth = process.stdout.columns ? process.stdout.columns - 10 : 120;
+      const head = summary.table[0];
+      const table = new Table({
+        head,
+        colWidths: Array(head.length).fill(Math.floor(maxWidth / head.length)),
+        wordWrap: true,
+        wrapOnWordBoundary: true,
+        style: {
+          head: ['blue', 'bold'],
+        },
+      });
+      // Skip first row (header) and add the rest. Color the first column green if it's a success, red if it's a failure.
+      for (const row of summary.table.slice(1)) {
+        const color = row[0] === 'PASS' ? 'green' : row[0].startsWith('FAIL') ? 'red' : undefined;
+        table.push(row.map((col, i) => (i === 0 && color ? chalk[color](col) : col)));
+      }
+
+      logger.info('\n' + table.toString());
+    }
+    logger.info('Evaluation complete');
+    logger.info(chalk.green.bold(`Successes: ${summary.stats.successes}`));
+    logger.info(chalk.red.bold(`Failures: ${summary.stats.failures}`));
+    logger.info(
+      `Token usage: Total ${summary.stats.tokenUsage.total} Prompt ${summary.stats.tokenUsage.prompt} Completion ${summary.stats.tokenUsage.completion}`,
+    );
+    logger.info('Done.');
+  });
+
+program.parse(process.argv);
+
+if (!process.argv.slice(2).length) {
+  program.outputHelp();
+}

package/src/providers.ts

@@ -0,0 +1,170 @@
+import fetch from 'node-fetch';
+
+import { ApiProvider, ProviderResponse } from './types.js';
+import logger from './logger.js';
+
+export class OpenAiGenericProvider implements ApiProvider {
+  modelName: string;
+  apiKey: string;
+
+  constructor(modelName: string, apiKey?: string) {
+    this.modelName = modelName;
+
+    const key = apiKey || process.env.OPENAI_API_KEY;
+    if (!key) {
+      throw new Error(
+        'OpenAI API key is not set. Set OPENAI_API_KEY environment variable or pass it as an argument to the constructor.',
+      );
+    }
+    this.apiKey = key;
+  }
+
+  id(): string {
+    return `openai:${this.modelName}`;
+  }
+
+  toString(): string {
+    return `[OpenAI Provider ${this.modelName}]`;
+  }
+
+  // @ts-ignore: Prompt is not used in this implementation
+  async callApi(prompt: string): Promise<ProviderResponse> {
+    throw new Error('Not implemented');
+  }
+}
+
+export class OpenAiCompletionProvider extends OpenAiGenericProvider {
+  static OPENAI_COMPLETION_MODELS = [
+    'text-davinci-003',
+    'text-davinci-002',
+    'text-curie-001',
+    'text-babbage-001',
+    'text-ada-001',
+  ];
+
+  constructor(modelName: string, apiKey?: string) {
+    if (!OpenAiCompletionProvider.OPENAI_COMPLETION_MODELS.includes(modelName)) {
+      throw new Error(
+        `Unknown OpenAI completion model name: ${modelName}. Use one of the following: ${OpenAiCompletionProvider.OPENAI_COMPLETION_MODELS.join(
+          ', ',
+        )}`,
+      );
+    }
+    super(modelName, apiKey);
+  }
+
+  async callApi(prompt: string): Promise<ProviderResponse> {
+    const body = {
+      model: this.modelName,
+      prompt,
+      max_tokens: process.env.OPENAI_MAX_TOKENS || 1024,
+      temperature: process.env.OPENAI_TEMPERATURE || 0,
+    };
+    logger.debug(`Calling OpenAI API: ${JSON.stringify(body)}`);
+    const response = await fetch('https://api.openai.com/v1/completions', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${this.apiKey}`,
+      },
+      body: JSON.stringify(body),
+    });
+
+    const data = (await response.json()) as unknown as any;
+    logger.debug(`\tOpenAI API response: ${JSON.stringify(data)}`);
+    return {
+      output: data.choices[0].text,
+      tokenUsage: {
+        total: data.usage.total_tokens,
+        prompt: data.usage.prompt_tokens,
+        completion: data.usage.completion_tokens,
+      },
+    };
+  }
+}
+
+export class OpenAiChatCompletionProvider extends OpenAiGenericProvider {
+  static OPENAI_CHAT_MODELS = [
+    'gpt-4',
+    'gpt-4-0314',
+    'gpt-4-32k',
+    'gpt-4-32k-0314',
+    'gpt-3.5-turbo',
+    'gpt-3.5-turbo-0301',
+  ];
+
+  constructor(modelName: string, apiKey?: string) {
+    if (!OpenAiChatCompletionProvider.OPENAI_CHAT_MODELS.includes(modelName)) {
+      throw new Error(
+        `Unknown OpenAI completion model name: ${modelName}. Use one of the following: ${OpenAiCompletionProvider.OPENAI_COMPLETION_MODELS.join(
+          ', ',
+        )}`,
+      );
+    }
+    super(modelName, apiKey);
+  }
+
+  async callApi(prompt: string): Promise<ProviderResponse> {
+    let messages: { role: string; content: string }[];
+    try {
+      // User can specify `messages` payload as JSON, or we'll just put the
+      // string prompt into a `messages` array.
+      messages = JSON.parse(prompt);
+    } catch (e) {
+      messages = [{ role: 'user', content: prompt }];
+    }
+    const body = {
+      model: this.modelName,
+      messages: messages,
+      max_tokens: process.env.OPENAI_MAX_TOKENS || 1024,
+      temperature: process.env.OPENAI_MAX_TEMPERATURE || 0,
+    };
+    logger.debug(`Calling OpenAI API: ${JSON.stringify(body)}`);
+    const response = await fetch('https://api.openai.com/v1/chat/completions', {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${this.apiKey}`,
+      },
+      body: JSON.stringify(body),
+    });
+
+    const data = (await response.json()) as unknown as any;
+    logger.debug(`\tOpenAI API response: ${JSON.stringify(data)}`);
+    return {
+      output: data.choices[0].message.content,
+      tokenUsage: {
+        total: data.usage.total_tokens,
+        prompt: data.usage.prompt_tokens,
+        completion: data.usage.completion_tokens,
+      },
+    };
+  }
+}
+
+export function loadApiProvider(providerPath: string): ApiProvider {
+  if (providerPath?.startsWith('openai:')) {
+    // Load OpenAI module
+    const options = providerPath.split(':');
+    const modelType = options[1];
+    const modelName = options[2];
+
+    if (modelType === 'chat') {
+      return new OpenAiChatCompletionProvider(modelName || 'gpt-3.5-turbo');
+    } else if (modelType === 'completion') {
+      return new OpenAiCompletionProvider(modelName || 'text-davinci-003');
+    } else if (OpenAiChatCompletionProvider.OPENAI_CHAT_MODELS.includes(modelType)) {
+      return new OpenAiChatCompletionProvider(modelType);
+    } else if (OpenAiCompletionProvider.OPENAI_COMPLETION_MODELS.includes(modelType)) {
+      return new OpenAiCompletionProvider(modelType);
+    } else {
+      throw new Error(
+        `Unknown OpenAI model type: ${modelType}. Use one of the following providers: openai:chat:<model name>, openai:completion:<model name>`,
+      );
+    }
+  }
+
+  // Load custom module
+  const CustomApiProvider = require(providerPath).default;
+  return new CustomApiProvider();
+}

package/src/tableOutput.html

@@ -0,0 +1,55 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width" />
+    <title>Table Output</title>
+    <style>
+      body {
+        font-family: -apple-system, BlinkMacSystemFont, Segoe UI, Roboto, Helvetica, Arial,
+          sans-serif;
+      }
+      table,
+      th,
+      td {
+        border: 1px solid black;
+        border-collapse: collapse;
+        text-align: left;
+        word-break: break-all;
+      }
+      th,
+      td {
+        padding: 5px;
+      }
+      /* If data-content is exactly "PASS", set font color to green */
+      tr > td[data-content='PASS']:first-child {
+        color: green;
+        font-weight: bold;
+      }
+
+      /* If data-content starts with "FAIL", set font color to red */
+      tr > td[data-content^='FAIL']:first-child {
+        color: red;
+        font-weight: bold;
+      }
+    </style>
+  </head>
+  <body>
+    <table>
+      <thead>
+        {% for header in table[0] %}
+        <th>{{ header }}</th>
+        {% endfor %}
+      </thead>
+      <tbody>
+        {% for row in table.slice(1) %}
+        <tr>
+          {% for cell in row %}
+          <td data-content="{{cell}}">{{ cell }}</td>
+          {% endfor %}
+        </tr>
+        {% endfor %}
+      </tbody>
+    </table>
+  </body>
+</html>

package/src/types.ts

@@ -0,0 +1,63 @@
+export interface CommandLineOptions {
+  prompt: string[];
+  provider: string[];
+  output?: string;
+  vars?: string;
+  config?: string;
+  verbose?: boolean;
+  maxConcurrency?: number;
+}
+
+export interface ApiProvider {
+  id: () => string;
+  callApi: (prompt: string) => Promise<ProviderResponse>;
+}
+
+interface TokenUsage {
+  total: number;
+  prompt: number;
+  completion: number;
+}
+
+export interface ProviderResponse {
+  output: string;
+  tokenUsage?: TokenUsage;
+}
+
+export interface CsvRow {
+  [key: string]: string;
+}
+
+export type VarMapping = Record<string, string>;
+
+export interface EvaluateOptions {
+  providers: ApiProvider[];
+  prompts: string[];
+  vars?: VarMapping[];
+
+  maxConcurrency?: number;
+  showProgressBar?: boolean;
+}
+
+export interface Prompt {
+  raw: string;
+  display: string;
+}
+
+export interface EvaluateResult {
+  prompt: Prompt;
+  vars: Record<string, string>;
+  response?: ProviderResponse;
+  error?: string;
+  success: boolean;
+}
+
+export interface EvaluateSummary {
+  results: EvaluateResult[];
+  table: string[][];
+  stats: {
+    successes: number;
+    failures: number;
+    tokenUsage: TokenUsage;
+  };
+}