ai-functions 2.1.3 → 2.4.0

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,
+      }
+    },
+  }
+}

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'