ai-functions 0.3.0 → 0.4.0

package/src/context.ts

@@ -0,0 +1,332 @@
+/**
+ * Execution Context for AI Functions
+ *
+ * Provides configuration without polluting function signatures.
+ * Settings flow from environment → global context → local context.
+ *
+ * @example
+ * ```ts
+ * // Set global defaults (from environment or initialization)
+ * configure({
+ *   provider: 'anthropic',
+ *   model: 'claude-sonnet-4-20250514',
+ *   batchMode: 'auto', // 'auto' | 'immediate' | 'deferred'
+ * })
+ *
+ * // Or use execution context for specific operations
+ * await withContext({ provider: 'openai', model: 'gpt-4o' }, async () => {
+ *   const titles = await list`10 blog titles`
+ *   return titles.map(title => write`blog post: ${title}`)
+ * })
+ * ```
+ *
+ * @packageDocumentation
+ */
+
+import type { FunctionOptions } from './template.js'
+import type { BatchProvider } from './batch-queue.js'
+
+// ============================================================================
+// Types
+// ============================================================================
+
+/** Batch execution mode */
+export type BatchMode =
+  | 'auto'       // Smart selection: immediate < flexThreshold, flex < batchThreshold, batch above
+  | 'immediate'  // Execute immediately (concurrent requests, full price)
+  | 'flex'       // Use flex processing (faster than batch, ~50% discount, minutes)
+  | 'deferred'   // Always use provider batch API (50% discount, up to 24hr)
+
+/** Execution context configuration */
+export interface ExecutionContext extends FunctionOptions {
+  /** Batch provider to use */
+  provider?: BatchProvider
+  /** Batch execution mode */
+  batchMode?: BatchMode
+  /** Minimum items to use flex processing (for 'auto' mode, default: 5) */
+  flexThreshold?: number
+  /** Minimum items to use batch API (for 'auto' mode, default: 500) */
+  batchThreshold?: number
+  /** Webhook URL for batch completion notifications */
+  webhookUrl?: string
+  /** Custom metadata for batch jobs */
+  metadata?: Record<string, unknown>
+}
+
+// ============================================================================
+// Global Context
+// ============================================================================
+
+let globalContext: ExecutionContext = {}
+
+/**
+ * Configure global defaults for AI functions
+ *
+ * @example
+ * ```ts
+ * configure({
+ *   model: 'claude-sonnet-4-20250514',
+ *   provider: 'anthropic',
+ *   batchMode: 'auto',
+ *   batchThreshold: 5,
+ * })
+ * ```
+ */
+export function configure(context: ExecutionContext): void {
+  globalContext = { ...globalContext, ...context }
+}
+
+/**
+ * Get the current global context
+ */
+export function getGlobalContext(): ExecutionContext {
+  return { ...globalContext }
+}
+
+/**
+ * Reset global context to defaults
+ */
+export function resetContext(): void {
+  globalContext = {}
+}
+
+// ============================================================================
+// Async Local Storage for Execution Context
+// ============================================================================
+
+// Use AsyncLocalStorage if available (Node.js), otherwise fallback to global
+let asyncLocalStorage: {
+  getStore: () => ExecutionContext | undefined
+  run: <T>(store: ExecutionContext, callback: () => T) => T
+} | null = null
+
+// Lazy initialization of AsyncLocalStorage
+let asyncLocalStorageInitialized = false
+
+async function initAsyncLocalStorage(): Promise<void> {
+  if (asyncLocalStorageInitialized) return
+  asyncLocalStorageInitialized = true
+
+  try {
+    const { AsyncLocalStorage } = await import('async_hooks')
+    asyncLocalStorage = new AsyncLocalStorage<ExecutionContext>()
+  } catch {
+    // Not in Node.js environment, use global context only
+  }
+}
+
+// Initialize synchronously if possible (for Node.js environments)
+if (typeof process !== 'undefined' && process.versions?.node) {
+  import('async_hooks')
+    .then(({ AsyncLocalStorage }) => {
+      asyncLocalStorage = new AsyncLocalStorage<ExecutionContext>()
+      asyncLocalStorageInitialized = true
+    })
+    .catch(() => {
+      asyncLocalStorageInitialized = true
+    })
+}
+
+/**
+ * Get the current execution context
+ * Merges: environment defaults → global context → local context
+ */
+export function getContext(): ExecutionContext {
+  const envContext = getEnvContext()
+  const localContext = asyncLocalStorage?.getStore()
+
+  return {
+    ...envContext,
+    ...globalContext,
+    ...localContext,
+  }
+}
+
+/**
+ * Run a function with a specific execution context
+ *
+ * @example
+ * ```ts
+ * const posts = await withContext({ provider: 'openai', batchMode: 'deferred' }, async () => {
+ *   const titles = await list`10 blog titles`
+ *   return titles.map(title => write`blog post: ${title}`)
+ * })
+ * ```
+ */
+export function withContext<T>(
+  context: ExecutionContext,
+  fn: () => T | Promise<T>
+): T | Promise<T> {
+  const mergedContext = { ...getContext(), ...context }
+
+  if (asyncLocalStorage) {
+    return asyncLocalStorage.run(mergedContext, fn)
+  }
+
+  // Fallback: temporarily modify global context
+  const previousContext = globalContext
+  globalContext = mergedContext
+  try {
+    return fn()
+  } finally {
+    globalContext = previousContext
+  }
+}
+
+// ============================================================================
+// Environment Defaults
+// ============================================================================
+
+function getEnvContext(): ExecutionContext {
+  if (typeof process === 'undefined') return {}
+
+  const context: ExecutionContext = {}
+
+  // Model defaults
+  if (process.env.AI_MODEL) {
+    context.model = process.env.AI_MODEL
+  }
+
+  // Provider defaults
+  if (process.env.AI_PROVIDER) {
+    context.provider = process.env.AI_PROVIDER as BatchProvider
+  } else if (process.env.ANTHROPIC_API_KEY && !process.env.OPENAI_API_KEY) {
+    context.provider = 'anthropic'
+  } else if (process.env.OPENAI_API_KEY) {
+    context.provider = 'openai'
+  } else if (process.env.CLOUDFLARE_API_TOKEN) {
+    context.provider = 'cloudflare'
+  } else if (process.env.AWS_ACCESS_KEY_ID) {
+    context.provider = 'bedrock'
+  }
+
+  // Batch mode
+  if (process.env.AI_BATCH_MODE) {
+    context.batchMode = process.env.AI_BATCH_MODE as BatchMode
+  }
+
+  // Flex threshold (when to start using flex processing)
+  if (process.env.AI_FLEX_THRESHOLD) {
+    context.flexThreshold = parseInt(process.env.AI_FLEX_THRESHOLD, 10)
+  }
+
+  // Batch threshold (when to switch from flex to full batch)
+  if (process.env.AI_BATCH_THRESHOLD) {
+    context.batchThreshold = parseInt(process.env.AI_BATCH_THRESHOLD, 10)
+  }
+
+  // Webhook URL
+  if (process.env.AI_BATCH_WEBHOOK_URL) {
+    context.webhookUrl = process.env.AI_BATCH_WEBHOOK_URL
+  }
+
+  return context
+}
+
+// ============================================================================
+// Context Helpers
+// ============================================================================
+
+/**
+ * Get the effective model from context
+ */
+export function getModel(): string {
+  const ctx = getContext()
+  return ctx.model || 'sonnet'
+}
+
+/**
+ * Get the effective provider from context
+ */
+export function getProvider(): BatchProvider {
+  const ctx = getContext()
+  return ctx.provider || 'openai'
+}
+
+/**
+ * Get the effective batch mode from context
+ */
+export function getBatchMode(): BatchMode {
+  const ctx = getContext()
+  return ctx.batchMode || 'auto'
+}
+
+/**
+ * Get the flex threshold from context (minimum items to use flex)
+ * Default: 5 items
+ */
+export function getFlexThreshold(): number {
+  const ctx = getContext()
+  return ctx.flexThreshold || 5
+}
+
+/**
+ * Get the batch threshold from context (minimum items to use full batch)
+ * Default: 500 items
+ */
+export function getBatchThreshold(): number {
+  const ctx = getContext()
+  return ctx.batchThreshold || 500
+}
+
+/** Execution tier for processing */
+export type ExecutionTier = 'immediate' | 'flex' | 'batch'
+
+/**
+ * Determine the execution tier for a given number of items
+ *
+ * Auto mode tiers:
+ * - immediate: < flexThreshold (default 5) - concurrent requests, full price
+ * - flex: flexThreshold to batchThreshold (5-500) - ~50% discount, minutes
+ * - batch: >= batchThreshold (500+) - 50% discount, up to 24hr
+ *
+ * @example
+ * ```ts
+ * getExecutionTier(3)   // 'immediate' (< 5)
+ * getExecutionTier(50)  // 'flex' (5-500)
+ * getExecutionTier(1000) // 'batch' (500+)
+ * ```
+ */
+export function getExecutionTier(itemCount: number): ExecutionTier {
+  const mode = getBatchMode()
+
+  switch (mode) {
+    case 'immediate':
+      return 'immediate'
+    case 'flex':
+      return 'flex'
+    case 'deferred':
+      return 'batch'
+    case 'auto':
+    default: {
+      const flexThreshold = getFlexThreshold()
+      const batchThreshold = getBatchThreshold()
+
+      if (itemCount < flexThreshold) {
+        return 'immediate'
+      } else if (itemCount < batchThreshold) {
+        return 'flex'
+      } else {
+        return 'batch'
+      }
+    }
+  }
+}
+
+/**
+ * Check if we should use the batch API for a given number of items
+ * @deprecated Use getExecutionTier() instead for more granular control
+ */
+export function shouldUseBatchAPI(itemCount: number): boolean {
+  const tier = getExecutionTier(itemCount)
+  return tier === 'flex' || tier === 'batch'
+}
+
+/**
+ * Check if flex processing is available for the current provider
+ * Only OpenAI and AWS Bedrock support flex processing currently
+ */
+export function isFlexAvailable(): boolean {
+  const provider = getProvider()
+  return provider === 'openai' || provider === 'bedrock' || provider === 'google'
+}

package/src/embeddings.ts

@@ -0,0 +1,244 @@
+/**
+ * Embedding utilities from AI SDK
+ *
+ * Re-exports embed, embedMany, and cosineSimilarity from the Vercel AI SDK
+ * with additional convenience wrappers.
+ *
+ * Default model: Cloudflare Workers AI @cf/baai/bge-m3
+ *
+ * @packageDocumentation
+ */
+
+// Re-export core embedding functions from AI SDK
+export { embed, embedMany, cosineSimilarity } from 'ai'
+
+// Re-export types
+export type {
+  EmbeddingModel,
+  Embedding
+} from 'ai'
+
+// Re-export Cloudflare provider from ai-providers
+export { cloudflare, cloudflareEmbedding, DEFAULT_CF_EMBEDDING_MODEL } from 'ai-providers/cloudflare'
+
+import { embed as aiEmbed, embedMany as aiEmbedMany } from 'ai'
+import { cloudflareEmbedding, DEFAULT_CF_EMBEDDING_MODEL } from 'ai-providers/cloudflare'
+
+/**
+ * Get the default embedding model (Cloudflare @cf/baai/bge-m3)
+ */
+export function getDefaultEmbeddingModel() {
+  return cloudflareEmbedding(DEFAULT_CF_EMBEDDING_MODEL)
+}
+
+/**
+ * Embed a single value using the default Cloudflare model
+ *
+ * @example
+ * ```ts
+ * import { embedText } from 'ai-functions'
+ *
+ * const { embedding } = await embedText('hello world')
+ * ```
+ */
+export async function embedText(value: string) {
+  return aiEmbed({
+    model: getDefaultEmbeddingModel(),
+    value
+  })
+}
+
+/**
+ * Embed multiple values using the default Cloudflare model
+ *
+ * @example
+ * ```ts
+ * import { embedTexts } from 'ai-functions'
+ *
+ * const { embeddings } = await embedTexts(['doc1', 'doc2', 'doc3'])
+ * ```
+ */
+export async function embedTexts(values: string[]) {
+  return aiEmbedMany({
+    model: getDefaultEmbeddingModel(),
+    values
+  })
+}
+
+/**
+ * Result of an embed operation
+ */
+export interface EmbedResult<T = string> {
+  /** The original input value */
+  value: T
+  /** The generated embedding vector */
+  embedding: number[]
+  /** Token usage */
+  usage: {
+    tokens: number
+  }
+}
+
+/**
+ * Result of an embedMany operation
+ */
+export interface EmbedManyResult<T = string> {
+  /** The original input values */
+  values: T[]
+  /** The generated embedding vectors */
+  embeddings: number[][]
+  /** Token usage */
+  usage: {
+    tokens: number
+  }
+}
+
+/**
+ * Find the most similar items to a query embedding
+ *
+ * @example
+ * ```ts
+ * import { embed, embedMany, findSimilar } from 'ai-functions'
+ *
+ * const documents = ['doc1', 'doc2', 'doc3']
+ * const { embeddings } = await embedMany({ model, values: documents })
+ * const { embedding: queryEmbedding } = await embed({ model, value: 'search query' })
+ *
+ * const results = findSimilar(queryEmbedding, embeddings, documents, { topK: 2 })
+ * // [{ item: 'doc1', score: 0.95, index: 0 }, { item: 'doc2', score: 0.82, index: 1 }]
+ * ```
+ */
+export function findSimilar<T>(
+  queryEmbedding: number[],
+  embeddings: number[][],
+  items: T[],
+  options: {
+    /** Number of results to return (default: 10) */
+    topK?: number
+    /** Minimum similarity score (default: 0) */
+    minScore?: number
+  } = {}
+): Array<{ item: T; score: number; index: number }> {
+  const { topK = 10, minScore = 0 } = options
+
+  // Import cosineSimilarity dynamically to avoid issues if ai isn't installed
+  const { cosineSimilarity } = require('ai')
+
+  const scored = embeddings
+    .map((embedding, index) => ({
+      item: items[index]!,
+      score: cosineSimilarity(queryEmbedding, embedding) as number,
+      index
+    }))
+    .filter(result => result.score >= minScore)
+    .sort((a, b) => b.score - a.score)
+    .slice(0, topK)
+
+  return scored
+}
+
+/**
+ * Calculate pairwise similarities between all embeddings
+ *
+ * @example
+ * ```ts
+ * const matrix = pairwiseSimilarity(embeddings)
+ * // matrix[i][j] = similarity between embeddings[i] and embeddings[j]
+ * ```
+ */
+export function pairwiseSimilarity(embeddings: number[][]): number[][] {
+  const { cosineSimilarity } = require('ai')
+
+  const n = embeddings.length
+  const matrix: number[][] = Array(n).fill(null).map(() => Array(n).fill(0))
+
+  for (let i = 0; i < n; i++) {
+    matrix[i]![i] = 1 // Self-similarity is always 1
+    for (let j = i + 1; j < n; j++) {
+      const sim = cosineSimilarity(embeddings[i], embeddings[j])
+      matrix[i]![j] = sim
+      matrix[j]![i] = sim
+    }
+  }
+
+  return matrix
+}
+
+/**
+ * Cluster embeddings by similarity using a simple threshold-based approach
+ *
+ * @example
+ * ```ts
+ * const clusters = clusterBySimilarity(embeddings, items, { threshold: 0.8 })
+ * // [[item1, item2], [item3], [item4, item5, item6]]
+ * ```
+ */
+export function clusterBySimilarity<T>(
+  embeddings: number[][],
+  items: T[],
+  options: {
+    /** Similarity threshold for clustering (default: 0.8) */
+    threshold?: number
+  } = {}
+): T[][] {
+  const { threshold = 0.8 } = options
+  const { cosineSimilarity } = require('ai')
+
+  const n = embeddings.length
+  const assigned = new Set<number>()
+  const clusters: T[][] = []
+
+  for (let i = 0; i < n; i++) {
+    if (assigned.has(i)) continue
+
+    const cluster: T[] = [items[i]!]
+    assigned.add(i)
+
+    for (let j = i + 1; j < n; j++) {
+      if (assigned.has(j)) continue
+
+      const sim = cosineSimilarity(embeddings[i], embeddings[j])
+      if (sim >= threshold) {
+        cluster.push(items[j]!)
+        assigned.add(j)
+      }
+    }
+
+    clusters.push(cluster)
+  }
+
+  return clusters
+}
+
+/**
+ * Average multiple embeddings into a single embedding
+ * Useful for creating document embeddings from chunk embeddings
+ */
+export function averageEmbeddings(embeddings: number[][]): number[] {
+  if (embeddings.length === 0) return []
+
+  const dim = embeddings[0]!.length
+  const result = new Array(dim).fill(0)
+
+  for (const embedding of embeddings) {
+    for (let i = 0; i < dim; i++) {
+      result[i] += embedding[i]
+    }
+  }
+
+  const n = embeddings.length
+  for (let i = 0; i < dim; i++) {
+    result[i] /= n
+  }
+
+  return result
+}
+
+/**
+ * Normalize an embedding to unit length
+ */
+export function normalizeEmbedding(embedding: number[]): number[] {
+  const magnitude = Math.sqrt(embedding.reduce((sum, val) => sum + val * val, 0))
+  if (magnitude === 0) return embedding
+  return embedding.map(val => val / magnitude)
+}

package/src/eval/index.ts

@@ -0,0 +1,8 @@
+/**
+ * AI Functions Eval Suite
+ *
+ * @packageDocumentation
+ */
+
+export * from './models.js'
+export * from './runner.js'