ai-functions 2.1.1 → 2.3.0

package/src/context.ts

@@ -1,8 +1,8 @@
 /**
  * Execution Context for AI Functions
  *
- * Provides configuration without polluting function signatures.
- * Settings flow from environment → global context → local context.
+ * Standalone context module with batch processing and budget tracking features.
+ * Settings flow from environment -> global context -> local context.
  *
  * @example
  * ```ts
@@ -23,24 +23,60 @@
  * @packageDocumentation
  */
 
-import type { FunctionOptions } from './template.js'
 import type { BatchProvider } from './batch-queue.js'
+import type { RequestContext as IRequestContext, ModelPricing } from './budget.js'
 
 // ============================================================================
-// Types
+// Core Types (from template.ts equivalent)
+// ============================================================================
+
+/**
+ * Common options for all AI functions
+ */
+export interface FunctionOptions {
+  /** Model to use (e.g., 'claude-opus-4-5', 'gpt-5-1', 'gemini-3-pro') */
+  model?: string
+  /** Thinking level: 'none', 'low', 'medium', 'high', or token budget number */
+  thinking?: 'none' | 'low' | 'medium' | 'high' | number
+  /** Temperature (0-2) */
+  temperature?: number
+  /** Maximum tokens to generate */
+  maxTokens?: number
+  /** System prompt */
+  system?: string
+  /** Processing mode */
+  mode?: 'default' | 'background'
+}
+
+// ============================================================================
+// Extended Types for Batch Processing
 // ============================================================================
 
 /** 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)
+  | '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)
+
+/** Budget configuration for context */
+export interface ContextBudgetConfig {
+  /** Maximum total tokens allowed */
+  maxTokens?: number
+  /** Maximum cost in USD */
+  maxCost?: number
+  /** Alert thresholds as fractions (e.g., [0.5, 0.8, 1.0]) */
+  alertThresholds?: number[]
+  /** Custom pricing for models not in default pricing table */
+  customPricing?: Record<string, ModelPricing>
+}
 
-/** Execution context configuration */
+/**
+ * Execution context with batch processing and budget features.
+ */
 export interface ExecutionContext extends FunctionOptions {
-  /** Batch provider to use */
-  provider?: BatchProvider
+  /** Provider to use (for model resolution) */
+  provider?: BatchProvider | string
   /** Batch execution mode */
   batchMode?: BatchMode
   /** Minimum items to use flex processing (for 'auto' mode, default: 5) */
@@ -51,6 +87,10 @@ export interface ExecutionContext extends FunctionOptions {
   webhookUrl?: string
   /** Custom metadata for batch jobs */
   metadata?: Record<string, unknown>
+  /** Budget configuration for tracking and limits */
+  budget?: ContextBudgetConfig
+  /** Request context for tracing */
+  requestContext?: IRequestContext
 }
 
 // ============================================================================
@@ -103,18 +143,6 @@ let asyncLocalStorage: {
 // 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')
@@ -127,16 +155,80 @@ if (typeof process !== 'undefined' && process.versions?.node) {
     })
 }
 
+// ============================================================================
+// 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']
+  } 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'
+  }
+
+  return context
+}
+
+function getEnvBatchContext(): Partial<ExecutionContext> {
+  if (typeof process === 'undefined') return {}
+
+  const context: Partial<ExecutionContext> = {}
+
+  // 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 Functions
+// ============================================================================
+
 /**
  * Get the current execution context
- * Merges: environment defaults → global context → local context
+ * Merges: environment defaults -> global context -> local context
  */
 export function getContext(): ExecutionContext {
   const envContext = getEnvContext()
+  const envBatchContext = getEnvBatchContext()
   const localContext = asyncLocalStorage?.getStore()
 
   return {
     ...envContext,
+    ...envBatchContext,
     ...globalContext,
     ...localContext,
   }
@@ -173,56 +265,6 @@ export function withContext<T>(
   }
 }
 
-// ============================================================================
-// 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
 // ============================================================================
@@ -236,13 +278,17 @@ export function getModel(): string {
 }
 
 /**
- * Get the effective provider from context
+ * Get the effective provider from context (typed as BatchProvider)
  */
 export function getProvider(): BatchProvider {
   const ctx = getContext()
-  return ctx.provider || 'openai'
+  return (ctx.provider || 'openai') as BatchProvider
 }
 
+// ============================================================================
+// Batch-Specific Context Helpers
+// ============================================================================
+
 /**
  * Get the effective batch mode from context
  */