feat: Add BRAINTRUST_CACHE_LOCATION env var to control caching location

.changeset/itchy-falcons-flow.md

@@ -0,0 +1,5 @@
+---
+"braintrust": minor
+---
+
+feat: Add `BRAINTRUST_CACHE_LOCATION` env var to control caching location

js/src/logger.ts

@@ -148,8 +148,7 @@ import {
   devNullWritableStream,
 } from "./functions/stream";
 import iso, { IsoAsyncLocalStorage } from "./isomorph";
-import { canUseDiskCache, DiskCache } from "./prompt-cache/disk-cache";
-import { LRUCache } from "./prompt-cache/lru-cache";
+import { createCacheLayers } from "./prompt-cache/cache-config";
 import { PromptCache } from "./prompt-cache/prompt-cache";
 import { ParametersCache } from "./prompt-cache/parameters-cache";
 import {
@@ -711,34 +710,25 @@ export class BraintrustState {
 
     this.resetLoginInfo();
 
-    const memoryCache = new LRUCache<string, Prompt>({
-      max: Number(iso.getEnv("BRAINTRUST_PROMPT_CACHE_MEMORY_MAX")) ?? 1 << 10,
+    const { memoryCache, diskCache } = createCacheLayers<Prompt>({
+      memoryMaxEnvVar: "BRAINTRUST_PROMPT_CACHE_MEMORY_MAX",
+      diskCacheDirEnvVar: "BRAINTRUST_PROMPT_CACHE_DIR",
+      diskMaxEnvVar: "BRAINTRUST_PROMPT_CACHE_DISK_MAX",
+      getDefaultDiskCacheDir: () =>
+        `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/prompt_cache`,
     });
-    const diskCache = canUseDiskCache()
-      ? new DiskCache<Prompt>({
-          cacheDir:
-            iso.getEnv("BRAINTRUST_PROMPT_CACHE_DIR") ??
-            `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/prompt_cache`,
-          max:
-            Number(iso.getEnv("BRAINTRUST_PROMPT_CACHE_DISK_MAX")) ?? 1 << 20,
-        })
-      : undefined;
     this.promptCache = new PromptCache({ memoryCache, diskCache });
 
-    const parametersMemoryCache = new LRUCache<string, RemoteEvalParameters>({
-      max:
-        Number(iso.getEnv("BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX")) ?? 1 << 10,
+    const {
+      memoryCache: parametersMemoryCache,
+      diskCache: parametersDiskCache,
+    } = createCacheLayers<RemoteEvalParameters>({
+      memoryMaxEnvVar: "BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX",
+      diskCacheDirEnvVar: "BRAINTRUST_PARAMETERS_CACHE_DIR",
+      diskMaxEnvVar: "BRAINTRUST_PARAMETERS_CACHE_DISK_MAX",
+      getDefaultDiskCacheDir: () =>
+        `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/parameters_cache`,
     });
-    const parametersDiskCache = canUseDiskCache()
-      ? new DiskCache<RemoteEvalParameters>({
-          cacheDir:
-            iso.getEnv("BRAINTRUST_PARAMETERS_CACHE_DIR") ??
-            `${iso.getEnv("HOME") ?? iso.homedir!()}/.braintrust/parameters_cache`,
-          max:
-            Number(iso.getEnv("BRAINTRUST_PARAMETERS_CACHE_DISK_MAX")) ??
-            1 << 20,
-        })
-      : undefined;
     this.parametersCache = new ParametersCache({
       memoryCache: parametersMemoryCache,
       diskCache: parametersDiskCache,

js/src/prompt-cache/cache-config.ts

@@ -0,0 +1,99 @@
+import { debugLogger } from "../debug-logger";
+import iso from "../isomorph";
+import { canUseDiskCache, DiskCache } from "./disk-cache";
+import { LRUCache } from "./lru-cache";
+
+type CacheMode = "mixed" | "memory" | "disk" | "none";
+
+const CACHE_LOCATION_ENV_VAR = "BRAINTRUST_CACHE_LOCATION";
+// Cache max values are entry counts, not byte sizes.
+const DEFAULT_CACHE_MEMORY_MAX = 1 << 10; // 2^10 = 1024 entries.
+const DEFAULT_CACHE_DISK_MAX = 1 << 20; // 2^20 = 1,048,576 entries.
+let warnedInvalidCacheModeEnvValue = false;
+let warnedUnavailableDiskCacheMode = false;
+
+function warnInvalidCacheMode(value: string) {
+  if (warnedInvalidCacheModeEnvValue) {
+    return;
+  }
+  warnedInvalidCacheModeEnvValue = true;
+  debugLogger.warn(
+    `Invalid ${CACHE_LOCATION_ENV_VAR} value "${value}". Expected "mixed", "memory", "disk", or "none". Falling back to "mixed".`,
+  );
+}
+
+function warnUnavailableDiskCache() {
+  if (warnedUnavailableDiskCacheMode) {
+    return;
+  }
+  warnedUnavailableDiskCacheMode = true;
+  debugLogger.warn(
+    `Disk cache is not supported on this platform, so ${CACHE_LOCATION_ENV_VAR}="disk" disables prompt and parameters caching.`,
+  );
+}
+
+function parseCacheMode(): CacheMode {
+  const value = iso.getEnv(CACHE_LOCATION_ENV_VAR);
+  const normalized = value?.trim().toLowerCase();
+  if (!normalized) {
+    return "mixed";
+  }
+  if (
+    normalized === "mixed" ||
+    normalized === "memory" ||
+    normalized === "disk" ||
+    normalized === "none"
+  ) {
+    return normalized;
+  }
+  warnInvalidCacheMode(value ?? "");
+  return "mixed";
+}
+
+function parsePositiveIntegerEnv(envVar: string, defaultValue: number): number {
+  const value = Number(iso.getEnv(envVar));
+  return Number.isInteger(value) && value > 0 ? value : defaultValue;
+}
+
+export function createCacheLayers<T>({
+  memoryMaxEnvVar,
+  diskCacheDirEnvVar,
+  diskMaxEnvVar,
+  getDefaultDiskCacheDir,
+}: {
+  memoryMaxEnvVar: string;
+  diskCacheDirEnvVar: string;
+  diskMaxEnvVar: string;
+  getDefaultDiskCacheDir: () => string;
+}): {
+  memoryCache?: LRUCache<string, T>;
+  diskCache?: DiskCache<T>;
+} {
+  const mode = parseCacheMode();
+  const memoryCache =
+    mode === "mixed" || mode === "memory"
+      ? new LRUCache<string, T>({
+          max: parsePositiveIntegerEnv(
+            memoryMaxEnvVar,
+            DEFAULT_CACHE_MEMORY_MAX,
+          ),
+        })
+      : undefined;
+
+  let diskCache: DiskCache<T> | undefined;
+  if (mode === "mixed" || mode === "disk") {
+    if (canUseDiskCache()) {
+      diskCache = new DiskCache<T>({
+        cacheDir: iso.getEnv(diskCacheDirEnvVar) ?? getDefaultDiskCacheDir(),
+        max: parsePositiveIntegerEnv(diskMaxEnvVar, DEFAULT_CACHE_DISK_MAX),
+      });
+    } else if (mode === "disk") {
+      warnUnavailableDiskCache();
+    }
+  }
+
+  if (diskCache) {
+    return { memoryCache, diskCache };
+  }
+  return { memoryCache };
+}