promptfoo

1
import path from 'path';
2
import fs from 'fs';
3

4
import logger from '../logger';
5
import { getCache, isCacheEnabled } from '../cache';
6
import { runPython } from '../python/wrapper';
7
import { sha256 } from '../util';
8

9
import type {
10
  ApiProvider,
11
  CallApiContextParams,
12
  ProviderOptions,
13
  ProviderResponse,
14
} from '../types';
15

16
interface PythonProviderConfig {
17
  pythonExecutable?: string;
18
}
19

20
export class PythonProvider implements ApiProvider {
21
  private config: PythonProviderConfig;
22

23
  constructor(private scriptPath: string, private options?: ProviderOptions) {
24
    this.id = () => options?.id ?? `python:${this.scriptPath}`;
25
    this.config = options?.config ?? {};
26
  }
27

28
  id() {
29
    return `python:${this.scriptPath}`;
30
  }
31

32
  async callApi(prompt: string, context?: CallApiContextParams): Promise<ProviderResponse> {
33
    const absPath = path.resolve(path.join(this.options?.config.basePath || '', this.scriptPath));
34
    logger.debug(`Computing file hash for script ${absPath}`);
35
    const fileHash = sha256(fs.readFileSync(absPath, 'utf-8'));
36
    const cacheKey = `python:${this.scriptPath}:${fileHash}:${prompt}:${JSON.stringify(
37
      this.options,
38
    )}`;
39
    const cache = await getCache();
40
    let cachedResult;
41

42
    if (isCacheEnabled()) {
43
      cachedResult = (await cache.get(cacheKey)) as string;
44
    }
45

46
    if (cachedResult) {
47
      logger.debug(`Returning cached result for script ${absPath}`);
48
      return JSON.parse(cachedResult);
49
    } else {
50
      const args = [prompt, this.options, context];
51
      logger.debug(
52
        `Running python script ${absPath} with scriptPath ${this.scriptPath} and args: ${args.join(
53
          '\n',
54
        )}`,
55
      );
56
      const result = (await runPython(absPath, 'call_api', args, {
57
        pythonExecutable: this.config.pythonExecutable,
58
      })) as {
59
        output?: string;
60
        error?: string;
61
      };
62
      if (!('output' in result) && !('error' in result)) {
63
        throw new Error(
64
          'The Python script `call_api` function must return a dict with an `output` or `error` string',
65
        );
66
      }
67
      if (isCacheEnabled()) {
68
        await cache.set(cacheKey, JSON.stringify(result));
69
      }
70
      return result;
71
    }
72
  }
73
}
74