dd-trace 6.15.0 → 6.17.0

package/packages/dd-trace/src/llmobs/prompts/cache.js

@@ -0,0 +1,272 @@
+'use strict'
+
+const { createHash, randomUUID } = require('node:crypto')
+const fs = require('node:fs')
+const os = require('node:os')
+const path = require('node:path')
+
+const { LRUCache } = require('../../../../../vendor/dist/lru-cache')
+
+const log = require('../../log')
+const ManagedPrompt = require('./prompt')
+
+const CACHE_SUBPATH = path.join('datadog', 'llmobs', 'prompts')
+const MAX_HOT_ENTRIES = 1024
+const PROMPT_SOURCES = new Set(['registry', 'cache', 'fallback', 'ff', 'resolve'])
+
+/** @param {string} key */
+function promptIdFromKey (key) {
+  return key.slice(0, key.lastIndexOf(':'))
+}
+
+/**
+ * @param {string} promptId
+ * @param {Array<unknown>} selector
+ */
+function cacheKey (promptId, selector) {
+  const hash = createHash('sha1').update(JSON.stringify(selector)).digest('hex').slice(0, 16)
+  return `${promptId}:${hash}`
+}
+
+function defaultCacheDir () {
+  try {
+    const home = os.homedir()
+    if (home) return path.join(home, '.cache', CACHE_SUBPATH)
+  } catch {}
+  return path.join(os.tmpdir(), CACHE_SUBPATH)
+}
+
+class HotCache {
+  /**
+   * Create the bounded in-memory prompt cache.
+   * @param {object} options
+   * @param {number} options.ttlMs
+   * @param {(key: string, stale: ManagedPrompt | undefined, options: {context: object}) =>
+   *   Promise<ManagedPrompt | undefined>} [options.fetchMethod]
+   */
+  constructor ({ ttlMs, fetchMethod }) {
+    this.enabled = ttlMs > 0
+    this.cache = new LRUCache({
+      max: MAX_HOT_ENTRIES,
+      ttl: this.enabled && Number.isFinite(ttlMs) ? ttlMs : undefined,
+      allowStale: true,
+      noDeleteOnStaleGet: true,
+      fetchMethod,
+    })
+  }
+
+  /**
+   * Read a prompt and its freshness.
+   * @param {string} key
+   */
+  get (key) {
+    if (!this.enabled) return
+    const status = {}
+    const prompt = this.cache.get(key, { status })
+    if (!prompt) return
+    return { prompt, stale: status.get === 'stale' }
+  }
+
+  /**
+   * Store a prompt.
+   * @param {string} key
+   * @param {ManagedPrompt} prompt
+   * @param {number} [ageMs]
+   */
+  set (key, prompt, ageMs = 0) {
+    if (this.enabled) this.cache.set(key, prompt, { start: performance.now() - ageMs })
+  }
+
+  /**
+   * Refresh a stale prompt in the background.
+   * @param {string} key
+   * @param {object} context
+   */
+  refresh (key, context) {
+    if (!this.enabled) return
+    void this.cache.fetch(key, {
+      allowStaleOnFetchRejection: true,
+      context,
+      forceRefresh: true,
+      noDeleteOnFetchRejection: true,
+    }).catch(() => {})
+  }
+
+  /**
+   * Delete one selector.
+   * @param {string} key
+   */
+  delete (key) {
+    if (this.enabled) this.cache.delete(key)
+  }
+
+  /**
+   * Clear all hot prompt entries.
+   */
+  clear () {
+    if (this.enabled) this.cache.clear()
+  }
+
+  /**
+   * Evict every selector for one exact prompt ID.
+   * @param {string} promptId
+   */
+  evictPrompt (promptId) {
+    if (!this.enabled) return
+    for (const [key] of this.cache.dump()) {
+      if (promptIdFromKey(key) === promptId) this.cache.delete(key)
+    }
+  }
+}
+
+class WarmCache {
+  /**
+   * Create the optional filesystem prompt cache.
+   * @param {object} options
+   * @param {string | undefined} options.cacheDir
+   * @param {boolean} [options.enabled]
+   * @param {number} options.ttlMs
+   */
+  constructor ({ cacheDir, enabled = true, ttlMs }) {
+    this.enabled = enabled && ttlMs > 0
+    this.ttlMs = ttlMs
+    this.cacheDir = cacheDir || defaultCacheDir()
+    if (this.enabled) this.#ensureDir(this.cacheDir)
+  }
+
+  /**
+   * Create a secure cache directory.
+   * @param {string} directory
+   */
+  #ensureDir (directory) {
+    try {
+      fs.mkdirSync(directory, { recursive: true, mode: 0o700 })
+    } catch (error) {
+      log.warn('Failed to create prompt cache directory: %s', error.message)
+      this.enabled = false
+    }
+  }
+
+  /**
+   * Resolve the collision-safe directory for a prompt ID.
+   * @param {string} promptId
+   */
+  #promptDir (promptId) {
+    return path.join(this.cacheDir, Buffer.from(promptId).toString('base64url') || '_')
+  }
+
+  /**
+   * Resolve the cache file for a selector key.
+   * @param {string} key
+   */
+  #path (key) {
+    const index = key.lastIndexOf(':')
+    return path.join(this.#promptDir(key.slice(0, index)), `${key.slice(index + 1)}.json`)
+  }
+
+  /**
+   * Restore and validate a prompt from the warm cache.
+   * @param {object} data
+   */
+  #deserialize (data) {
+    const validTemplate = typeof data?.template === 'string' || (
+      Array.isArray(data?.template) && data.template.every(message => {
+        return message && typeof message.role === 'string' && typeof message.content === 'string'
+      })
+    )
+    if (
+      typeof data?.id !== 'string' ||
+      typeof data.version !== 'string' ||
+      !PROMPT_SOURCES.has(data.source) ||
+      !validTemplate ||
+      (data.promptUuid !== undefined && typeof data.promptUuid !== 'string') ||
+      (data.promptVersionUuid !== undefined && typeof data.promptVersionUuid !== 'string')
+    ) {
+      throw new TypeError('Invalid managed prompt cache entry')
+    }
+    return new ManagedPrompt(data)
+  }
+
+  /**
+   * Read a prompt and its freshness.
+   * @param {string} key
+   */
+  get (key) {
+    if (!this.enabled) return
+    try {
+      const data = JSON.parse(fs.readFileSync(this.#path(key), 'utf8'))
+      if (!Number.isFinite(data.timestamp)) throw new TypeError('Invalid prompt cache timestamp')
+      const prompt = this.#deserialize(data.prompt)
+      const ageMs = Math.max(0, Date.now() - data.timestamp)
+      return { prompt, stale: ageMs > this.ttlMs, ageMs }
+    } catch (error) {
+      log.debug('Failed to read prompt from cache: %s', error.message)
+    }
+  }
+
+  /**
+   * Store a prompt with restrictive permissions.
+   * @param {string} key
+   * @param {ManagedPrompt} prompt
+   */
+  set (key, prompt) {
+    if (!this.enabled) return
+    const file = this.#path(key)
+    const temporary = `${file}.tmp.${randomUUID()}`
+    try {
+      this.#ensureDir(path.dirname(file))
+      if (!this.enabled) return
+      // Write separately, then rename so other processes see either the old file or the complete new one.
+      fs.writeFileSync(temporary, JSON.stringify({ prompt, timestamp: Date.now() }), {
+        encoding: 'utf8',
+        mode: 0o600,
+      })
+      fs.renameSync(temporary, file)
+    } catch (error) {
+      try { fs.rmSync(temporary, { force: true }) } catch {}
+      log.debug('Failed to write prompt to cache: %s', error.message)
+    }
+  }
+
+  /**
+   * Delete one selector.
+   * @param {string} key
+   */
+  delete (key) {
+    if (!this.enabled) return
+    try {
+      fs.rmSync(this.#path(key), { force: true })
+    } catch (error) {
+      log.debug('Failed to delete prompt from cache: %s', error.message)
+    }
+  }
+
+  /**
+   * Evict every selector for one exact prompt ID.
+   * @param {string} promptId
+   */
+  evictPrompt (promptId) {
+    if (!this.enabled) return
+    try {
+      fs.rmSync(this.#promptDir(promptId), { recursive: true, force: true })
+    } catch (error) {
+      log.debug('Failed to evict prompt from cache: %s', error.message)
+    }
+  }
+
+  /**
+   * Clear all warm prompt entries.
+   */
+  clear () {
+    if (!this.enabled) return
+    try {
+      for (const entry of fs.readdirSync(this.cacheDir)) {
+        fs.rmSync(path.join(this.cacheDir, entry), { recursive: true, force: true })
+      }
+    } catch (error) {
+      log.debug('Failed to clear prompt cache: %s', error.message)
+    }
+  }
+}
+
+module.exports = { HotCache, WarmCache, cacheKey, promptIdFromKey }