ai-functions 2.1.1 → 2.3.0

package/src/logger.ts

@@ -0,0 +1,232 @@
+/**
+ * Configurable Logger for AI Functions
+ *
+ * Provides a pluggable logging interface that defaults to console but can be
+ * configured to use any logging implementation. This allows library consumers
+ * to integrate with their own logging infrastructure.
+ *
+ * @example
+ * ```ts
+ * import { configureLogger, getLogger } from 'ai-functions'
+ *
+ * // Use with default console logger
+ * const logger = getLogger()
+ * logger.warn('Something happened')
+ *
+ * // Configure a custom logger
+ * configureLogger({
+ *   debug: (msg, ...args) => myLogger.debug(msg, ...args),
+ *   info: (msg, ...args) => myLogger.info(msg, ...args),
+ *   warn: (msg, ...args) => myLogger.warn(msg, ...args),
+ *   error: (msg, ...args) => myLogger.error(msg, ...args),
+ * })
+ *
+ * // Disable logging entirely
+ * configureLogger(null)
+ * ```
+ *
+ * @packageDocumentation
+ */
+
+// ============================================================================
+// Logger Interface
+// ============================================================================
+
+/**
+ * Logger interface that matches the standard console methods.
+ *
+ * All methods are optional to allow partial implementations.
+ * Missing methods will be no-ops when called.
+ */
+export interface Logger {
+  /** Debug level logging (verbose) */
+  debug?: (message: string, ...args: unknown[]) => void
+  /** Info level logging (general information) */
+  info?: (message: string, ...args: unknown[]) => void
+  /** Warning level logging (potential issues) */
+  warn?: (message: string, ...args: unknown[]) => void
+  /** Error level logging (errors and exceptions) */
+  error?: (message: string, ...args: unknown[]) => void
+}
+
+/**
+ * Logger configuration options
+ */
+export interface LoggerConfig {
+  /** The logger implementation to use */
+  logger?: Logger | null
+  /** Minimum log level to output */
+  level?: 'debug' | 'info' | 'warn' | 'error' | 'silent'
+}
+
+// ============================================================================
+// Logger State
+// ============================================================================
+
+/** Log level priority (lower = more verbose) */
+const LOG_LEVELS: Record<string, number> = {
+  debug: 0,
+  info: 1,
+  warn: 2,
+  error: 3,
+  silent: 4,
+}
+
+/** Default console-based logger */
+const defaultLogger: Logger = {
+  debug: (message: string, ...args: unknown[]) => console.debug(message, ...args),
+  info: (message: string, ...args: unknown[]) => console.info(message, ...args),
+  warn: (message: string, ...args: unknown[]) => console.warn(message, ...args),
+  error: (message: string, ...args: unknown[]) => console.error(message, ...args),
+}
+
+/** No-op logger for when logging is disabled */
+const noopLogger: Logger = {
+  debug: () => {},
+  info: () => {},
+  warn: () => {},
+  error: () => {},
+}
+
+/** Current logger instance */
+let currentLogger: Logger = defaultLogger
+
+/** Current minimum log level */
+let currentLevel: string = 'warn'
+
+// ============================================================================
+// Configuration Functions
+// ============================================================================
+
+/**
+ * Configure the global logger for ai-functions.
+ *
+ * @param config - Logger configuration or just a logger instance
+ *
+ * @example
+ * ```ts
+ * // Use a custom logger
+ * configureLogger({
+ *   logger: {
+ *     debug: (msg) => myLogger.debug(msg),
+ *     info: (msg) => myLogger.info(msg),
+ *     warn: (msg) => myLogger.warn(msg),
+ *     error: (msg) => myLogger.error(msg),
+ *   },
+ *   level: 'info',
+ * })
+ *
+ * // Just set the log level
+ * configureLogger({ level: 'debug' })
+ *
+ * // Disable all logging
+ * configureLogger({ logger: null })
+ * // or
+ * configureLogger({ level: 'silent' })
+ *
+ * // Reset to defaults
+ * configureLogger({})
+ * ```
+ */
+export function configureLogger(config: LoggerConfig | Logger | null): void {
+  // Handle null (disable logging)
+  if (config === null) {
+    currentLogger = noopLogger
+    return
+  }
+
+  // Handle Logger directly
+  if (config && ('debug' in config || 'info' in config || 'warn' in config || 'error' in config)) {
+    // Check if it's a Logger (has at least one log method but no 'logger' or 'level' key)
+    if (!('logger' in config) && !('level' in config)) {
+      currentLogger = config as Logger
+      return
+    }
+  }
+
+  // Handle LoggerConfig
+  const loggerConfig = config as LoggerConfig
+
+  if (loggerConfig.logger === null) {
+    currentLogger = noopLogger
+  } else if (loggerConfig.logger) {
+    currentLogger = loggerConfig.logger
+  } else if (config === undefined || Object.keys(config as object).length === 0) {
+    // Reset to default
+    currentLogger = defaultLogger
+    currentLevel = 'warn'
+    return
+  }
+
+  if (loggerConfig.level) {
+    currentLevel = loggerConfig.level
+  }
+}
+
+/**
+ * Get the current logger instance.
+ *
+ * Returns a logger that respects the configured log level.
+ * Methods for levels below the configured minimum will be no-ops.
+ *
+ * @returns Logger instance with level filtering applied
+ *
+ * @example
+ * ```ts
+ * const logger = getLogger()
+ * logger.warn('This is a warning')
+ * logger.error('This is an error', { details: 'foo' })
+ * ```
+ */
+export function getLogger(): Required<Logger> {
+  const minLevel = LOG_LEVELS[currentLevel] ?? LOG_LEVELS['warn']!
+
+  return {
+    debug: (message: string, ...args: unknown[]) => {
+      if (minLevel <= LOG_LEVELS['debug']!) {
+        currentLogger.debug?.(message, ...args)
+      }
+    },
+    info: (message: string, ...args: unknown[]) => {
+      if (minLevel <= LOG_LEVELS['info']!) {
+        currentLogger.info?.(message, ...args)
+      }
+    },
+    warn: (message: string, ...args: unknown[]) => {
+      if (minLevel <= LOG_LEVELS['warn']!) {
+        currentLogger.warn?.(message, ...args)
+      }
+    },
+    error: (message: string, ...args: unknown[]) => {
+      if (minLevel <= LOG_LEVELS['error']!) {
+        currentLogger.error?.(message, ...args)
+      }
+    },
+  }
+}
+
+/**
+ * Reset the logger to default configuration.
+ *
+ * Restores the console-based logger with 'warn' minimum level.
+ */
+export function resetLogger(): void {
+  currentLogger = defaultLogger
+  currentLevel = 'warn'
+}
+
+/**
+ * Get the current log level
+ */
+export function getLogLevel(): string {
+  return currentLevel
+}
+
+/**
+ * Set the minimum log level
+ *
+ * @param level - Minimum level to log ('debug' | 'info' | 'warn' | 'error' | 'silent')
+ */
+export function setLogLevel(level: 'debug' | 'info' | 'warn' | 'error' | 'silent'): void {
+  currentLevel = level
+}

package/src/middleware/budget.ts

@@ -0,0 +1,171 @@
+/**
+ * budgetMiddleware — record token usage + cost into a {@link BudgetTracker}
+ *
+ * Replaces the post-hoc duck-typing in
+ * `services-as-software/src/v3/invoke/cost-estimate.ts` with a single
+ * AI-SDK-6 middleware: on `doGenerate` / `doStream` completion, read the
+ * `LanguageModelV3Usage` shape directly off the result and call
+ * `tracker.recordUsage(...)`. The pricing overlay is supplied via
+ * `customPricing` on the {@link BudgetTracker} (or we hand the tracker the
+ * pricing at construction time when the caller wants per-call isolation).
+ *
+ * Key V3 → BudgetTracker mapping detail: AI SDK 6 reports
+ * `usage.inputTokens.total` / `usage.outputTokens.total` as
+ * `number | undefined`. We coerce undefined → 0 so partial-streaming results
+ * (where the upstream provider didn't emit token counts) don't blow up the
+ * tracker. The `inputTokens.cacheRead` / `inputTokens.cacheWrite` breakdown
+ * is *not* propagated yet — round 13+ work to add prompt-cache awareness to
+ * BudgetTracker.
+ *
+ * Composition note: install **after** cache (so a cache hit still records
+ * the cost — the wrapped result is the same regardless of which layer
+ * served it) and **before** trace (so the trace event sees the final
+ * computed cost via the tracker).
+ *
+ * @packageDocumentation
+ */
+
+import type {
+  LanguageModelV3GenerateResult,
+  LanguageModelV3Middleware,
+  LanguageModelV3StreamPart,
+  LanguageModelV3StreamResult,
+  LanguageModelV3Usage,
+} from '@ai-sdk/provider'
+import type { BudgetTracker, ModelPricing } from '../budget.js'
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/**
+ * Pricing overlay supplied to the middleware. Mirrors the
+ * `BudgetConfig.customPricing` shape — keyed on model id, value is the
+ * per-million USD rate. Sourced (in services-as-software) from the
+ * `language-models/data/models.json` catalog so Llama / DeepSeek / Mistral /
+ * Qwen / Grok / Perplexity Sonar all get their real per-token rate.
+ */
+export type PricingOverlay = Record<string, ModelPricing>
+
+/** Options for {@link budgetMiddleware}. */
+export interface BudgetMiddlewareOptions {
+  /**
+   * The {@link BudgetTracker} to record usage into. Required — the
+   * middleware never constructs its own tracker (the tracker holds budget
+   * limits + alert callbacks, which the caller owns).
+   */
+  tracker: BudgetTracker
+  /**
+   * Pricing overlay (per-model rates). When supplied, takes precedence over
+   * the BudgetTracker's own default pricing for any matching model id. Pass
+   * the language-models catalog overlay here to extend pricing without
+   * mutating the tracker.
+   */
+  pricing?: PricingOverlay
+  /**
+   * Optional override for the model id reported to the tracker. Defaults to
+   * `model.modelId` (the wrapped model's underlying id). Pass an alias
+   * (`'sonnet'`, `'opus'`) to bridge to the alias-based pricing tables.
+   */
+  modelIdOverride?: string
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+function coerceUsage(usage: LanguageModelV3Usage | undefined): {
+  inputTokens: number
+  outputTokens: number
+} {
+  if (!usage) return { inputTokens: 0, outputTokens: 0 }
+  return {
+    inputTokens: usage.inputTokens?.total ?? 0,
+    outputTokens: usage.outputTokens?.total ?? 0,
+  }
+}
+
+function record(
+  tracker: BudgetTracker,
+  pricing: PricingOverlay | undefined,
+  modelId: string,
+  usage: LanguageModelV3Usage | undefined
+): void {
+  const { inputTokens, outputTokens } = coerceUsage(usage)
+  if (inputTokens === 0 && outputTokens === 0) return
+  // The pricing overlay is wired in via the tracker's `customPricing`
+  // already (set at BudgetTracker construction time by the caller). When
+  // the caller wants per-call pricing override, they install
+  // `pricing[modelId]` ahead of time. We expose `pricing` here as a
+  // forward-looking hook so we can later add per-call pricing without a
+  // breaking change.
+  void pricing
+  tracker.recordUsage({ inputTokens, outputTokens, model: modelId })
+}
+
+// ============================================================================
+// Middleware
+// ============================================================================
+
+/**
+ * Build a budget middleware for `wrapLanguageModel`. Records
+ * {@link LanguageModelV3Usage} into the supplied {@link BudgetTracker} on
+ * every successful `doGenerate` / `doStream` completion. Errors from the
+ * downstream model propagate unchanged — the tracker is only updated on
+ * success.
+ *
+ * For streaming calls, we accumulate the final `usage` from the `'finish'`
+ * stream part (per the V3 spec, the final `'finish'` event carries the
+ * authoritative usage shape) and record once on stream end.
+ *
+ * @example
+ * ```ts
+ * import { wrapLanguageModel } from 'ai'
+ * import { BudgetTracker, budgetMiddleware } from 'ai-functions'
+ *
+ * const tracker = new BudgetTracker({ maxCost: 1.0 })
+ * const model = wrapLanguageModel({
+ *   model: openai('gpt-4o'),
+ *   middleware: budgetMiddleware({ tracker }),
+ * })
+ * ```
+ */
+export function budgetMiddleware(options: BudgetMiddlewareOptions): LanguageModelV3Middleware {
+  const { tracker, pricing, modelIdOverride } = options
+  return {
+    specificationVersion: 'v3',
+    async wrapGenerate({ doGenerate, model }) {
+      const result = await doGenerate()
+      const modelId = modelIdOverride ?? model.modelId
+      record(tracker, pricing, modelId, result.usage)
+      return result
+    },
+    async wrapStream({ doStream, model }) {
+      const result = await doStream()
+      const modelId = modelIdOverride ?? model.modelId
+      let finalUsage: LanguageModelV3Usage | undefined
+      const transformedStream = result.stream.pipeThrough(
+        new TransformStream<LanguageModelV3StreamPart, LanguageModelV3StreamPart>({
+          transform(chunk, controller) {
+            if (chunk.type === 'finish') {
+              finalUsage = chunk.usage
+            }
+            controller.enqueue(chunk)
+          },
+          flush() {
+            record(tracker, pricing, modelId, finalUsage)
+          },
+        })
+      )
+      const wrapped: LanguageModelV3StreamResult = {
+        ...result,
+        stream: transformedStream,
+      }
+      return wrapped
+    },
+  }
+}
+
+// Re-export to make the result type available to consumers writing custom
+// middleware chains.
+export type { LanguageModelV3GenerateResult }

package/src/middleware/cache.ts

@@ -0,0 +1,299 @@
+/**
+ * cacheMiddleware — content-addressable cache for `wrapLanguageModel`
+ *
+ * Implements the AI SDK cookbook's local-caching-middleware pattern
+ * (https://ai-sdk.dev/cookbook/node/local-caching-middleware) on top of the
+ * AI SDK 6 `LanguageModelV3Middleware` shape:
+ *
+ * - **Hit derivation:** content-hash of `{ prompt, modelId, responseFormat }`
+ *   so a schema change (responseFormat.type === 'json' carries a `schema`
+ *   JSONSchema7) invalidates the entry. Generation parameters (temperature,
+ *   topP, etc.) are deliberately *not* part of the key for the eval-fixture
+ *   use case — flipping temperature shouldn't blow up a 5x verify-time win.
+ *   Callers who want strict keying should pass a custom `keyHash`.
+ *
+ * - **Stream support:** cached entries store the `LanguageModelV3StreamPart[]`
+ *   array; `wrapStream` replays them via `simulateReadableStream` so consumers
+ *   see the same chunked event sequence on a hit. (`wrapGenerate` is the
+ *   common path; both share the same cache map.)
+ *
+ * - **TTL:** 24h default, configurable via `ttlMs`. Entries past TTL are
+ *   evicted on access (lazy expiry — no background timer).
+ *
+ * - **Pluggable store:** in-memory default (Map-backed); `'disk'` writes to
+ *   a JSON file at `.cache/v3-eval-cache.json` for cross-process fixture
+ *   sharing. Disk reads/writes are best-effort — IO failures fall through
+ *   to the wrapped model.
+ *
+ * - **Env gate:** honors `process.env.V3_EVAL_CACHE`. When unset/empty, the
+ *   middleware short-circuits to a passthrough — useful for production where
+ *   cache hits would be incorrect but the operator wants the same wrap chain.
+ *   Set to `'1'` (or any truthy non-empty string) to enable.
+ *
+ * @packageDocumentation
+ */
+
+import { simulateReadableStream } from 'ai'
+import type {
+  LanguageModelV3CallOptions,
+  LanguageModelV3GenerateResult,
+  LanguageModelV3Middleware,
+  LanguageModelV3StreamPart,
+  LanguageModelV3StreamResult,
+} from '@ai-sdk/provider'
+import { hashKey } from '../cache.js'
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/** Cached payload — both generate result and stream chunks under one key. */
+interface CacheEntry {
+  /** Result captured from `doGenerate`. Absent if the entry came from a stream call. */
+  generateResult?: LanguageModelV3GenerateResult
+  /** Stream chunks captured from `doStream` (replayed via simulateReadableStream). */
+  streamChunks?: LanguageModelV3StreamPart[]
+  /** Insert epoch ms — drives TTL eviction. */
+  createdAt: number
+}
+
+/** Pluggable cache store for cached LLM results. */
+export interface CacheMiddlewareStore {
+  get(key: string): CacheEntry | undefined
+  set(key: string, value: CacheEntry): void
+  delete(key: string): void
+}
+
+/** Options for {@link cacheMiddleware}. */
+export interface CacheMiddlewareOptions {
+  /**
+   * Cache backend. `'memory'` uses a process-local Map; `'disk'` writes to
+   * `.cache/v3-eval-cache.json` for cross-process fixture sharing. A custom
+   * {@link CacheMiddlewareStore} can be passed instead.
+   *
+   * @default 'memory'
+   */
+  store?: 'memory' | 'disk' | CacheMiddlewareStore
+  /**
+   * TTL in milliseconds. Entries older than `ttlMs` are evicted on access.
+   *
+   * @default 86_400_000 (24h)
+   */
+  ttlMs?: number
+  /**
+   * Custom hash function for cache keys. Defaults to a stable hash of
+   * `{ prompt, modelId, responseFormat }`.
+   */
+  keyHash?: (params: LanguageModelV3CallOptions, modelId: string) => string
+  /**
+   * Optional override for the env gate. When `false`, the middleware acts
+   * as a passthrough regardless of `V3_EVAL_CACHE`. When `true`, always
+   * caches. Defaults to `process.env.V3_EVAL_CACHE` truthy-check.
+   */
+  enabled?: boolean
+  /** Optional custom path for the disk store (defaults to `.cache/v3-eval-cache.json`). */
+  diskPath?: string
+}
+
+// ============================================================================
+// Stores
+// ============================================================================
+
+class MemoryStore implements CacheMiddlewareStore {
+  private readonly map: Map<string, CacheEntry> = new Map()
+  get(key: string): CacheEntry | undefined {
+    return this.map.get(key)
+  }
+  set(key: string, value: CacheEntry): void {
+    this.map.set(key, value)
+  }
+  delete(key: string): void {
+    this.map.delete(key)
+  }
+}
+
+/**
+ * Disk-backed store. Best-effort — JSON parse / write errors fall through
+ * silently so a corrupt cache file never blocks an LLM call. The whole map
+ * is rewritten on each `set` (cheap for the eval-fixture use case which is
+ * dominated by reads).
+ */
+class DiskStore implements CacheMiddlewareStore {
+  private readonly path: string
+  private cache: Map<string, CacheEntry> | null = null
+
+  constructor(path: string) {
+    this.path = path
+  }
+
+  private load(): Map<string, CacheEntry> {
+    if (this.cache !== null) return this.cache
+    this.cache = new Map()
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-require-imports
+      const fs = require('fs') as typeof import('fs')
+      if (fs.existsSync(this.path)) {
+        const raw = fs.readFileSync(this.path, 'utf-8')
+        const parsed = JSON.parse(raw) as Record<string, CacheEntry>
+        for (const [k, v] of Object.entries(parsed)) {
+          this.cache.set(k, v)
+        }
+      }
+    } catch {
+      // best-effort
+    }
+    return this.cache
+  }
+
+  private flush(): void {
+    if (this.cache === null) return
+    try {
+      // eslint-disable-next-line @typescript-eslint/no-require-imports
+      const fs = require('fs') as typeof import('fs')
+      // eslint-disable-next-line @typescript-eslint/no-require-imports
+      const path = require('path') as typeof import('path')
+      const dir = path.dirname(this.path)
+      if (!fs.existsSync(dir)) {
+        fs.mkdirSync(dir, { recursive: true })
+      }
+      const obj = Object.fromEntries(this.cache)
+      fs.writeFileSync(this.path, JSON.stringify(obj), 'utf-8')
+    } catch {
+      // best-effort
+    }
+  }
+
+  get(key: string): CacheEntry | undefined {
+    return this.load().get(key)
+  }
+
+  set(key: string, value: CacheEntry): void {
+    this.load().set(key, value)
+    this.flush()
+  }
+
+  delete(key: string): void {
+    this.load().delete(key)
+    this.flush()
+  }
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+const DEFAULT_TTL_MS = 24 * 60 * 60 * 1000
+
+function defaultKeyHash(params: LanguageModelV3CallOptions, modelId: string): string {
+  // Stable hash of prompt + model + responseFormat (which carries the
+  // schema for object generation). Generation knobs are deliberately
+  // excluded so the eval-fixture cache survives temperature tweaks.
+  return hashKey({
+    prompt: params.prompt,
+    modelId,
+    responseFormat: params.responseFormat,
+  })
+}
+
+function envGateEnabled(): boolean {
+  const v = process.env['V3_EVAL_CACHE']
+  return typeof v === 'string' && v.length > 0
+}
+
+function isExpired(entry: CacheEntry, ttlMs: number): boolean {
+  return Date.now() - entry.createdAt > ttlMs
+}
+
+// ============================================================================
+// Middleware
+// ============================================================================
+
+/**
+ * Build a cache middleware for `wrapLanguageModel`. Wraps `doGenerate` and
+ * `doStream`; on a hit replays the cached payload, on a miss invokes the
+ * downstream model and stores the result.
+ *
+ * Composition note: install **before** budget/trace so cache hits don't
+ * pay the downstream model cost (the trace/budget middleware still see the
+ * payload via the wrapped result they observe in their own `wrapGenerate`).
+ *
+ * @example
+ * ```ts
+ * import { wrapLanguageModel } from 'ai'
+ * import { cacheMiddleware } from 'ai-functions'
+ *
+ * const model = wrapLanguageModel({
+ *   model: openai('gpt-4o'),
+ *   middleware: cacheMiddleware({ store: 'disk', ttlMs: 86_400_000 }),
+ * })
+ * ```
+ */
+export function cacheMiddleware(options: CacheMiddlewareOptions = {}): LanguageModelV3Middleware {
+  const ttlMs = options.ttlMs ?? DEFAULT_TTL_MS
+  const keyHash = options.keyHash ?? defaultKeyHash
+  const store: CacheMiddlewareStore =
+    options.store === undefined || options.store === 'memory'
+      ? new MemoryStore()
+      : options.store === 'disk'
+      ? new DiskStore(options.diskPath ?? '.cache/v3-eval-cache.json')
+      : options.store
+  const enabled = options.enabled ?? envGateEnabled()
+
+  return {
+    specificationVersion: 'v3',
+    async wrapGenerate({ doGenerate, params, model }) {
+      if (!enabled) return doGenerate()
+      const key = keyHash(params, model.modelId)
+      const cached = store.get(key)
+      if (cached !== undefined) {
+        if (isExpired(cached, ttlMs)) {
+          store.delete(key)
+        } else if (cached.generateResult !== undefined) {
+          return cached.generateResult
+        }
+      }
+      const result = await doGenerate()
+      store.set(key, { generateResult: result, createdAt: Date.now() })
+      return result
+    },
+    async wrapStream({ doStream, params, model }) {
+      if (!enabled) return doStream()
+      const key = keyHash(params, model.modelId)
+      const cached = store.get(key)
+      if (cached !== undefined) {
+        if (isExpired(cached, ttlMs)) {
+          store.delete(key)
+        } else if (cached.streamChunks !== undefined) {
+          // Replay cached chunks via simulateReadableStream so consumers
+          // see the same async iteration shape as a fresh call.
+          const replay: LanguageModelV3StreamResult = {
+            stream: simulateReadableStream<LanguageModelV3StreamPart>({
+              chunks: cached.streamChunks,
+              initialDelayInMs: 0,
+              chunkDelayInMs: 0,
+            }),
+          }
+          return replay
+        }
+      }
+      const result = await doStream()
+      // Tee the stream: forward to caller, accumulate for cache.
+      const chunks: LanguageModelV3StreamPart[] = []
+      const transformedStream = result.stream.pipeThrough(
+        new TransformStream<LanguageModelV3StreamPart, LanguageModelV3StreamPart>({
+          transform(chunk, controller) {
+            chunks.push(chunk)
+            controller.enqueue(chunk)
+          },
+          flush() {
+            store.set(key, { streamChunks: chunks, createdAt: Date.now() })
+          },
+        })
+      )
+      return {
+        ...result,
+        stream: transformedStream,
+      }
+    },
+  }
+}