ai-functions 2.1.1 → 2.3.0

package/src/middleware/embed-cache.ts

@@ -0,0 +1,195 @@
+/**
+ * embeddingCacheMiddleware — content-addressable cache for `wrapEmbeddingModel`
+ *
+ * Embedding-side analogue of {@link cacheMiddleware}. Wraps `doEmbed` and
+ * caches the resulting embeddings keyed on
+ * `{ values, modelId, providerOptions }` so a re-embed of the same value
+ * batch with the same model returns the cached vectors without hitting the
+ * provider.
+ *
+ * **Why a separate middleware instead of reusing `cacheMiddleware`?**
+ * AI SDK 6 splits language-model and embedding-model surfaces:
+ * `LanguageModelV3Middleware` exposes `wrapGenerate` / `wrapStream` against
+ * `LanguageModelV3CallOptions`, while `EmbeddingModelV3Middleware` exposes
+ * `wrapEmbed` against `EmbeddingModelV3CallOptions`. The cache shape
+ * (per-value vector vs. per-prompt completion payload) is also different —
+ * embeddings cache batched arrays, generations cache single result objects.
+ *
+ * - **Hit derivation:** stable hash of `{ values, modelId, providerOptions }`.
+ *   `values` is the array as-passed (caller can pre-normalise if they want
+ *   case/whitespace insensitivity). Generation knobs don't apply.
+ *
+ * - **Batch semantics:** the cache key is the *whole* batch. A subset hit
+ *   doesn't trigger a partial-fill — that's a more invasive shape change
+ *   (the legacy `EmbeddingCache.getMany` did per-text caching, but it was
+ *   only used in the example and added 100+ LOC of bookkeeping). Callers
+ *   that want per-text caching should use stable per-text batches.
+ *
+ * - **TTL:** 24h default, configurable. Lazy expiry on access.
+ *
+ * - **Pluggable store:** in-memory default (Map-backed); custom store
+ *   honored as-is. Disk persistence is intentionally not provided here —
+ *   embedding payloads (large `number[][]`) make on-disk JSON a bad fit;
+ *   callers who want it should pass a custom store.
+ *
+ * - **Env gate:** honors `process.env.V3_EVAL_CACHE` for parity with
+ *   `cacheMiddleware`. Override via the `enabled` option.
+ *
+ * @packageDocumentation
+ */
+
+import type {
+  EmbeddingModelV3CallOptions,
+  EmbeddingModelV3Embedding,
+  EmbeddingModelV3Middleware,
+  EmbeddingModelV3Result,
+  SharedV3Warning,
+} from '@ai-sdk/provider'
+import { hashKey } from '../cache.js'
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/** Cached embedding payload. */
+interface EmbedCacheEntry {
+  /** The embedding vectors returned for the cached batch. */
+  embeddings: Array<EmbeddingModelV3Embedding>
+  /** Provider warnings carried alongside the cached batch. */
+  warnings: Array<SharedV3Warning>
+  /** Insert epoch ms — drives TTL eviction. */
+  createdAt: number
+}
+
+/** Pluggable cache store for embedding results. */
+export interface EmbedCacheMiddlewareStore {
+  get(key: string): EmbedCacheEntry | undefined
+  set(key: string, value: EmbedCacheEntry): void
+  delete(key: string): void
+}
+
+/** Options for {@link embeddingCacheMiddleware}. */
+export interface EmbedCacheMiddlewareOptions {
+  /**
+   * Cache backend. `'memory'` uses a process-local Map. A custom
+   * {@link EmbedCacheMiddlewareStore} can be passed instead.
+   *
+   * @default 'memory'
+   */
+  store?: 'memory' | EmbedCacheMiddlewareStore
+  /**
+   * 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
+   * `{ values, modelId, providerOptions }`.
+   */
+  keyHash?: (params: EmbeddingModelV3CallOptions, 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
+}
+
+// ============================================================================
+// Stores
+// ============================================================================
+
+class MemoryStore implements EmbedCacheMiddlewareStore {
+  private readonly map: Map<string, EmbedCacheEntry> = new Map()
+  get(key: string): EmbedCacheEntry | undefined {
+    return this.map.get(key)
+  }
+  set(key: string, value: EmbedCacheEntry): void {
+    this.map.set(key, value)
+  }
+  delete(key: string): void {
+    this.map.delete(key)
+  }
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+const DEFAULT_TTL_MS = 24 * 60 * 60 * 1000
+
+function defaultKeyHash(params: EmbeddingModelV3CallOptions, modelId: string): string {
+  return hashKey({
+    values: params.values,
+    modelId,
+    providerOptions: params.providerOptions,
+  })
+}
+
+function envGateEnabled(): boolean {
+  const v = process.env['V3_EVAL_CACHE']
+  return typeof v === 'string' && v.length > 0
+}
+
+function isExpired(entry: EmbedCacheEntry, ttlMs: number): boolean {
+  return Date.now() - entry.createdAt > ttlMs
+}
+
+// ============================================================================
+// Middleware
+// ============================================================================
+
+/**
+ * Build an embedding-cache middleware for `wrapEmbeddingModel`.
+ *
+ * @example
+ * ```ts
+ * import { wrapEmbeddingModel } from 'ai'
+ * import { embeddingCacheMiddleware } from 'ai-functions'
+ *
+ * const model = wrapEmbeddingModel({
+ *   model: openai.embedding('text-embedding-3-small'),
+ *   middleware: embeddingCacheMiddleware({ ttlMs: 86_400_000 }),
+ * })
+ * ```
+ */
+export function embeddingCacheMiddleware(
+  options: EmbedCacheMiddlewareOptions = {}
+): EmbeddingModelV3Middleware {
+  const ttlMs = options.ttlMs ?? DEFAULT_TTL_MS
+  const keyHash = options.keyHash ?? defaultKeyHash
+  const store: EmbedCacheMiddlewareStore =
+    options.store === undefined || options.store === 'memory' ? new MemoryStore() : options.store
+  const enabled = options.enabled ?? envGateEnabled()
+
+  return {
+    specificationVersion: 'v3',
+    async wrapEmbed({ doEmbed, params, model }) {
+      if (!enabled) return doEmbed()
+      const key = keyHash(params, model.modelId)
+      const cached = store.get(key)
+      if (cached !== undefined) {
+        if (isExpired(cached, ttlMs)) {
+          store.delete(key)
+        } else {
+          // Replay shape matches EmbeddingModelV3Result. Provider-side
+          // metadata (response headers, body, usage) is intentionally absent
+          // on a hit — callers reading those should disable the cache.
+          const replay: EmbeddingModelV3Result = {
+            embeddings: cached.embeddings,
+            warnings: cached.warnings,
+          }
+          return replay
+        }
+      }
+      const result = await doEmbed()
+      store.set(key, {
+        embeddings: result.embeddings,
+        warnings: result.warnings,
+        createdAt: Date.now(),
+      })
+      return result
+    },
+  }
+}

package/src/middleware/index.ts

@@ -0,0 +1,23 @@
+/**
+ * Middleware barrel — composable AI SDK 6 `LanguageModelV3Middleware`
+ * primitives for `wrapLanguageModel`.
+ *
+ * @packageDocumentation
+ */
+
+export { cacheMiddleware, type CacheMiddlewareOptions, type CacheMiddlewareStore } from './cache.js'
+
+export {
+  embeddingCacheMiddleware,
+  type EmbedCacheMiddlewareOptions,
+  type EmbedCacheMiddlewareStore,
+} from './embed-cache.js'
+
+export { budgetMiddleware, type BudgetMiddlewareOptions, type PricingOverlay } from './budget.js'
+
+export {
+  traceMiddleware,
+  type TraceEvent,
+  type TraceEventKind,
+  type TraceMiddlewareOptions,
+} from './trace.js'

package/src/middleware/trace.ts

@@ -0,0 +1,248 @@
+/**
+ * traceMiddleware — emit per-call trace events for `wrapLanguageModel`
+ *
+ * Wraps `doGenerate` / `doStream` and emits a {@link TraceEvent} on every
+ * completion. The sink is opaque (caller supplies `emit`) so this primitive
+ * works equally well piping into:
+ *
+ *   - the v3 cascade-walker InvocationEvent stream (round 16+ work to add
+ *     `'persona-trace'` / `'cascade-trace'` to the union),
+ *   - an {@link import('../eval-log/index.js').EvalLogStore} for fixture
+ *     replay,
+ *   - OpenTelemetry / Datadog / Honeycomb adapters that map the event into
+ *     a span.
+ *
+ * **Emit-error tolerance:** if the supplied `emit` throws, we *swallow* the
+ * error (with a one-time `console.warn`) so a flaky trace sink can never
+ * break the wrapped LLM call. This matches the Evalite v0.19 trace
+ * middleware behaviour.
+ *
+ * Composition note: install **last** so the event sees the final outcome
+ * (post-cache, post-budget). The event's `costUsd` field is best-effort —
+ * the trace middleware doesn't have direct access to the budget tracker, so
+ * the caller can pass a `getCostUsd` resolver if they want costs in the
+ * event payload.
+ *
+ * @packageDocumentation
+ */
+
+import type {
+  LanguageModelV3CallOptions,
+  LanguageModelV3GenerateResult,
+  LanguageModelV3Middleware,
+  LanguageModelV3StreamPart,
+  LanguageModelV3StreamResult,
+  LanguageModelV3Usage,
+} from '@ai-sdk/provider'
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/**
+ * Discriminator for the originating call site. Callers inject this via the
+ * `kind` option so a single sink can fan events into different downstream
+ * streams (persona panel vs. cascade walker vs. ad-hoc test).
+ */
+export type TraceEventKind = 'persona-trace' | 'cascade-trace' | 'eval-trace' | string
+
+/**
+ * Trace event payload emitted on every wrapped call completion.
+ *
+ * Field design notes:
+ *   - `prompt` / `response` are stringified for cheap downstream storage
+ *     (the structured `LanguageModelV3Prompt` / `LanguageModelV3Content[]`
+ *     shapes are intentionally flattened).
+ *   - `usage` is the raw V3 shape (with the cache breakdown) — the
+ *     EvalLogStore consumer flattens it into total counts.
+ *   - `costUsd` is optional because the trace middleware doesn't compute
+ *     cost itself; callers either pass a resolver or compute downstream
+ *     from `usage`.
+ */
+export interface TraceEvent {
+  kind: TraceEventKind
+  model: string
+  prompt: string
+  response: string
+  usage: LanguageModelV3Usage | undefined
+  costUsd?: number
+  durationMs: number
+  /** Optional caller-supplied tags for downstream filtering. */
+  tags?: Record<string, string>
+}
+
+/** Options for {@link traceMiddleware}. */
+export interface TraceMiddlewareOptions {
+  /**
+   * Opaque sink. Errors thrown from `emit` are swallowed (with a one-time
+   * `console.warn`) so a flaky sink never breaks the wrapped LLM call.
+   */
+  emit: (event: TraceEvent) => void | Promise<void>
+  /**
+   * Discriminator threaded into the event's `kind` field. Defaults to
+   * `'eval-trace'`.
+   */
+  kind?: TraceEventKind
+  /**
+   * Optional cost resolver. When supplied, called with the V3 usage shape
+   * and the modelId; result is set on `event.costUsd`. Useful when the
+   * caller has a side-channel pricing table (the budgetMiddleware's
+   * tracker) and wants costs in the trace event itself.
+   */
+  getCostUsd?: (modelId: string, usage: LanguageModelV3Usage | undefined) => number
+  /** Optional caller-supplied tags merged into every emitted event. */
+  tags?: Record<string, string>
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+/**
+ * Flatten the structured V3 prompt into a single string for cheap storage.
+ * Walks system / user / assistant / tool messages and concatenates their
+ * text parts. Non-text parts (files, tool results) are summarised with a
+ * short marker so the trace doesn't grow unboundedly.
+ */
+function stringifyPrompt(params: LanguageModelV3CallOptions): string {
+  const out: string[] = []
+  for (const msg of params.prompt) {
+    if (msg.role === 'system') {
+      out.push(`[system] ${msg.content}`)
+      continue
+    }
+    if (typeof msg.content === 'string') {
+      out.push(`[${msg.role}] ${msg.content}`)
+      continue
+    }
+    if (Array.isArray(msg.content)) {
+      const parts: string[] = []
+      for (const part of msg.content) {
+        if (part.type === 'text') parts.push(part.text)
+        else parts.push(`[${part.type}]`)
+      }
+      out.push(`[${msg.role}] ${parts.join(' ')}`)
+    }
+  }
+  return out.join('\n')
+}
+
+/**
+ * Flatten the V3 generate result content into a single string. Walks the
+ * `content` array (text, reasoning, tool-call, etc.) and concatenates text
+ * parts; non-text parts get short summaries.
+ */
+function stringifyContent(content: LanguageModelV3GenerateResult['content']): string {
+  const parts: string[] = []
+  for (const part of content) {
+    if (part.type === 'text') parts.push(part.text)
+    else if (part.type === 'reasoning') parts.push(`[reasoning] ${part.text}`)
+    else parts.push(`[${part.type}]`)
+  }
+  return parts.join('')
+}
+
+let _hasWarnedEmit = false
+
+async function safeEmit(emit: TraceMiddlewareOptions['emit'], event: TraceEvent): Promise<void> {
+  try {
+    await emit(event)
+  } catch (err) {
+    if (!_hasWarnedEmit) {
+      _hasWarnedEmit = true
+      // eslint-disable-next-line no-console
+      console.warn(
+        `[ai-functions/traceMiddleware] emit() threw — subsequent emit errors will be silenced. ${
+          err instanceof Error ? err.message : String(err)
+        }`
+      )
+    }
+  }
+}
+
+// ============================================================================
+// Middleware
+// ============================================================================
+
+/**
+ * Build a trace middleware for `wrapLanguageModel`. Emits a
+ * {@link TraceEvent} on every successful `doGenerate` / `doStream`
+ * completion. Errors from `emit` are swallowed (one-time warn) so a flaky
+ * trace sink can never break the wrapped LLM call.
+ *
+ * @example
+ * ```ts
+ * import { wrapLanguageModel } from 'ai'
+ * import { traceMiddleware, getEvalLogStore } from 'ai-functions'
+ *
+ * const store = getEvalLogStore()
+ * const model = wrapLanguageModel({
+ *   model: openai('gpt-4o'),
+ *   middleware: traceMiddleware({
+ *     kind: 'cascade-trace',
+ *     emit: (event) => store.record({ ...event, costUsd: event.costUsd ?? 0 }),
+ *   }),
+ * })
+ * ```
+ */
+export function traceMiddleware(options: TraceMiddlewareOptions): LanguageModelV3Middleware {
+  const { emit, kind = 'eval-trace', getCostUsd, tags } = options
+  return {
+    specificationVersion: 'v3',
+    async wrapGenerate({ doGenerate, params, model }) {
+      const start = Date.now()
+      const result = await doGenerate()
+      const durationMs = Date.now() - start
+      const modelId = model.modelId
+      const event: TraceEvent = {
+        kind,
+        model: modelId,
+        prompt: stringifyPrompt(params),
+        response: stringifyContent(result.content),
+        usage: result.usage,
+        durationMs,
+        ...(getCostUsd !== undefined ? { costUsd: getCostUsd(modelId, result.usage) } : {}),
+        ...(tags !== undefined ? { tags } : {}),
+      }
+      await safeEmit(emit, event)
+      return result
+    },
+    async wrapStream({ doStream, params, model }) {
+      const start = Date.now()
+      const result = await doStream()
+      const modelId = model.modelId
+      let finalUsage: LanguageModelV3Usage | undefined
+      const collected: string[] = []
+      const transformedStream = result.stream.pipeThrough(
+        new TransformStream<LanguageModelV3StreamPart, LanguageModelV3StreamPart>({
+          transform(chunk, controller) {
+            if (chunk.type === 'text-delta') collected.push(chunk.delta)
+            else if (chunk.type === 'finish') finalUsage = chunk.usage
+            controller.enqueue(chunk)
+          },
+          flush() {
+            const durationMs = Date.now() - start
+            const event: TraceEvent = {
+              kind,
+              model: modelId,
+              prompt: stringifyPrompt(params),
+              response: collected.join(''),
+              usage: finalUsage,
+              durationMs,
+              ...(getCostUsd !== undefined ? { costUsd: getCostUsd(modelId, finalUsage) } : {}),
+              ...(tags !== undefined ? { tags } : {}),
+            }
+            // Fire-and-forget — TransformStream.flush is sync; we don't
+            // await safeEmit so a slow sink doesn't block stream close.
+            void safeEmit(emit, event)
+          },
+        })
+      )
+      const wrapped: LanguageModelV3StreamResult = {
+        ...result,
+        stream: transformedStream,
+      }
+      return wrapped
+    },
+  }
+}