promptfoo

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

4
import cacheManager from 'cache-manager';
5
import fsStore from 'cache-manager-fs-hash';
6

7
import logger from './logger';
8
import { fetchWithRetries } from './fetch';
9
import { getConfigDirectoryPath } from './util';
10

11
import type { Cache } from 'cache-manager';
12
import type { RequestInfo, RequestInit } from 'node-fetch';
13

14
let cacheInstance: Cache | undefined;
15

16
let enabled =
17
  typeof process.env.PROMPTFOO_CACHE_ENABLED === 'undefined'
18
    ? true
19
    : process.env.PROMPTFOO_CACHE_ENABLED === '1' ||
20
      process.env.PROMPTFOO_CACHE_ENABLED === 'true' ||
21
      process.env.PROMPTFOO_CACHE_ENABLED === 'yes';
22

23
const cacheType =
24
  process.env.PROMPTFOO_CACHE_TYPE || (process.env.NODE_ENV === 'test' ? 'memory' : 'disk');
25

26
export function getCache() {
27
  if (!cacheInstance) {
28
    const cachePath =
29
      process.env.PROMPTFOO_CACHE_PATH || path.join(getConfigDirectoryPath(), 'cache');
30
    if (!fs.existsSync(cachePath)) {
31
      logger.info(`Creating cache folder at ${cachePath}.`);
32
      fs.mkdirSync(cachePath, { recursive: true });
33
    }
34
    cacheInstance = cacheManager.caching({
35
      store: cacheType === 'disk' ? fsStore : 'memory',
36
      options: {
37
        max: process.env.PROMPTFOO_CACHE_MAX_FILE_COUNT || 10_000, // number of files
38
        path: cachePath,
39
        ttl: process.env.PROMPTFOO_CACHE_TTL || 60 * 60 * 24 * 14, // in seconds, 14 days
40
        maxsize: process.env.PROMPTFOO_CACHE_MAX_SIZE || 1e7, // in bytes, 10mb
41
        //zip: true, // whether to use gzip compression
42
      },
43
    });
44
  }
45
  return cacheInstance;
46
}
47

48
export async function fetchWithCache(
49
  url: RequestInfo,
50
  options: RequestInit = {},
51
  timeout: number,
52
  format: 'json' | 'text' = 'json',
53
  bust: boolean = false,
54
): Promise<{ data: any; cached: boolean }> {
55
  if (!enabled || bust) {
56
    const resp = await fetchWithRetries(url, options, timeout);
57
    const respText = await resp.text();
58
    try {
59
      return {
60
        cached: false,
61
        data: format === 'json' ? JSON.parse(respText) : respText,
62
      };
63
    } catch (error) {
64
      throw new Error(`Error parsing response as JSON: ${respText}`);
65
    }
66
  }
67

68
  const cache = await getCache();
69

70
  const copy = Object.assign({}, options);
71
  delete copy.headers;
72
  const cacheKey = `fetch:${url}:${JSON.stringify(copy)}`;
73

74
  // Try to get the cached response
75
  const cachedResponse = await cache.get(cacheKey);
76

77
  if (cachedResponse) {
78
    logger.debug(`Returning cached response for ${url}: ${cachedResponse}`);
79
    return {
80
      cached: true,
81
      data: JSON.parse(cachedResponse as string),
82
    };
83
  }
84

85
  // Fetch the actual data and store it in the cache
86
  const response = await fetchWithRetries(url, options, timeout);
87
  const responseText = await response.text();
88
  try {
89
    const data = format === 'json' ? JSON.parse(responseText) : responseText;
90
    if (response.ok) {
91
      logger.debug(`Storing ${url} response in cache: ${JSON.stringify(data)}`);
92
      await cache.set(cacheKey, JSON.stringify(data));
93
    }
94
    return {
95
      cached: false,
96
      data,
97
    };
98
  } catch (err) {
99
    throw new Error(
100
      `Error parsing response from ${url}: ${
101
        (err as Error).message
102
      }. Received text: ${responseText}`,
103
    );
104
  }
105
}
106

107
export function enableCache() {
108
  enabled = true;
109
}
110

111
export function disableCache() {
112
  enabled = false;
113
}
114

115
export async function clearCache() {
116
  return getCache().reset();
117
}
118

119
export function isCacheEnabled() {
120
  return enabled;
121
}
122