ai-functions 2.1.3 → 2.4.0

package/src/batch/provider.ts

@@ -0,0 +1,336 @@
+/**
+ * BatchProvider Port
+ *
+ * Defines the explicit port (interface) that every batch provider adapter must
+ * satisfy, plus helpers that concentrate logic genuinely shared across adapters
+ * (polling, concurrent processing, in-memory job tracking, JSON-Schema conversion).
+ *
+ * Each provider file (`anthropic.ts`, `openai.ts`, `google.ts`, `bedrock.ts`,
+ * `cloudflare.ts`, `memory.ts`) is now a small adapter that satisfies this port
+ * and concentrates on provider-specific HTTP/SDK calls and request/response shapes.
+ *
+ * Port + 4 adapters = real seam. The port lives here so provider files don't
+ * import shared logic from each other and can't accidentally diverge.
+ *
+ * @packageDocumentation
+ */
+
+import type {
+  BatchAdapter,
+  BatchItem,
+  BatchJob,
+  BatchQueueOptions,
+  BatchResult,
+  BatchStatus,
+} from '../batch-queue.js'
+
+// Re-export the port and its types so adapters can import everything from
+// `./provider.js` rather than reaching into `../batch-queue.js`.
+export type {
+  BatchAdapter,
+  BatchItem,
+  BatchJob,
+  BatchQueueOptions,
+  BatchResult,
+  BatchSubmitResult,
+  BatchProvider,
+  BatchStatus,
+  FlexAdapter,
+} from '../batch-queue.js'
+
+export { registerBatchAdapter, registerFlexAdapter } from '../batch-queue.js'
+
+// ============================================================================
+// Polling helper (shared by all adapters that have async batch jobs)
+// ============================================================================
+
+/** Terminal states for a batch job — polling stops here. */
+const TERMINAL_STATUSES: ReadonlySet<BatchStatus> = new Set([
+  'completed',
+  'failed',
+  'cancelled',
+  'expired',
+])
+
+/** Statuses that should trigger result fetch. */
+const RESULT_STATUSES: ReadonlySet<BatchStatus> = new Set(['completed', 'cancelled', 'failed'])
+
+/**
+ * Default `waitForCompletion` implementation built on top of `getStatus` +
+ * `getResults`. Adapters with non-standard completion semantics can still
+ * override `waitForCompletion`, but most don't need to.
+ *
+ * @param adapter   The batch adapter to poll
+ * @param batchId   The batch id to poll
+ * @param options.pollInterval  Poll interval in ms (default: 5000)
+ * @param options.fetchResultsOn  Statuses for which results should be fetched.
+ *                                Defaults to `completed`, `cancelled`, `failed`.
+ * @param options.throwOn  Statuses that should throw rather than fetch results.
+ *                         Useful for OpenAI which throws on `cancelled`/`expired`.
+ */
+export async function pollUntilComplete(
+  adapter: Pick<BatchAdapter, 'getStatus' | 'getResults'>,
+  batchId: string,
+  options: {
+    pollInterval?: number
+    fetchResultsOn?: ReadonlySet<BatchStatus>
+    throwOn?: ReadonlySet<BatchStatus>
+  } = {}
+): Promise<BatchResult[]> {
+  const pollInterval = options.pollInterval ?? 5000
+  const fetchResultsOn = options.fetchResultsOn ?? RESULT_STATUSES
+  const throwOn = options.throwOn
+
+  while (true) {
+    const status = await adapter.getStatus(batchId)
+
+    if (throwOn?.has(status.status)) {
+      throw new Error(`Batch ${status.status}`)
+    }
+
+    if (fetchResultsOn.has(status.status) || TERMINAL_STATUSES.has(status.status)) {
+      return adapter.getResults(batchId)
+    }
+
+    await sleep(pollInterval)
+  }
+}
+
+// ============================================================================
+// Concurrent processing helper (shared by flex adapters and "local" providers)
+// ============================================================================
+
+/**
+ * Run `processItem` over `items` with bounded concurrency, optionally
+ * sleeping between waves to respect provider rate limits. Per-item failures
+ * are caught and emitted as `{ status: 'failed', error }` results so a single
+ * bad item never poisons the whole batch.
+ *
+ * Used by flex adapters (OpenAI, Google, Bedrock) and by the "local" providers
+ * (Google, Bedrock, Cloudflare) that fake batch processing with concurrent
+ * direct API calls.
+ */
+export async function processConcurrently(
+  items: BatchItem[],
+  processItem: (item: BatchItem) => Promise<BatchResult>,
+  options: {
+    concurrency?: number
+    delayBetweenWaves?: number
+    onWaveComplete?: (results: BatchResult[]) => void
+  } = {}
+): Promise<BatchResult[]> {
+  const concurrency = options.concurrency ?? 10
+  const delay = options.delayBetweenWaves ?? 0
+  const results: BatchResult[] = []
+
+  for (let i = 0; i < items.length; i += concurrency) {
+    const wave = items.slice(i, i + concurrency)
+
+    const waveResults = await Promise.all(
+      wave.map(async (item) => {
+        try {
+          return await processItem(item)
+        } catch (error) {
+          return failedResult(item, error)
+        }
+      })
+    )
+
+    results.push(...waveResults)
+    options.onWaveComplete?.(results)
+
+    if (delay > 0 && i + concurrency < items.length) {
+      await sleep(delay)
+    }
+  }
+
+  return results
+}
+
+/** Build a `failed` BatchResult from an unknown thrown value. */
+export function failedResult(item: BatchItem, error: unknown): BatchResult {
+  return {
+    id: item.id,
+    customId: item.id,
+    status: 'failed',
+    error: error instanceof Error ? error.message : 'Unknown error',
+  }
+}
+
+// ============================================================================
+// LocalJobStore — in-memory job tracking shared by google/bedrock/cloudflare
+// ============================================================================
+
+/**
+ * Internal job state for adapters that don't have a real provider-side batch
+ * API and need to track jobs locally (Google, Bedrock, Cloudflare).
+ */
+export interface LocalJobState {
+  items: BatchItem[]
+  options: BatchQueueOptions
+  results: BatchResult[]
+  status: BatchStatus
+  createdAt: Date
+  completedAt?: Date
+  /** Optional adapter-specific metadata (e.g. Bedrock's jobArn) */
+  meta?: Record<string, unknown>
+}
+
+/**
+ * Per-provider in-memory job registry. Encapsulates the
+ * `Map<jobId, state>` + counter + status/result lookup pattern that
+ * google/bedrock/cloudflare were each duplicating.
+ */
+export class LocalJobStore {
+  private readonly jobs = new Map<string, LocalJobState>()
+  private counter = 0
+
+  constructor(private readonly idPrefix: string) {}
+
+  create(items: BatchItem[], options: BatchQueueOptions): { id: string; state: LocalJobState } {
+    const id = `${this.idPrefix}_${++this.counter}_${Date.now()}`
+    const state: LocalJobState = {
+      items,
+      options,
+      results: [],
+      status: 'pending',
+      createdAt: new Date(),
+    }
+    this.jobs.set(id, state)
+    return { id, state }
+  }
+
+  get(id: string): LocalJobState {
+    const state = this.jobs.get(id)
+    if (!state) {
+      throw new Error(`Batch not found: ${id}`)
+    }
+    return state
+  }
+
+  has(id: string): boolean {
+    return this.jobs.has(id)
+  }
+
+  /** Build a `BatchJob` snapshot for a tracked job. */
+  snapshot(id: string, provider: BatchJob['provider']): BatchJob {
+    const state = this.get(id)
+    const completedItems = state.results.filter((r) => r.status === 'completed').length
+    const failedItems = state.results.filter((r) => r.status === 'failed').length
+
+    return {
+      id,
+      provider,
+      status: state.status,
+      totalItems: state.items.length,
+      completedItems,
+      failedItems,
+      createdAt: state.createdAt,
+      ...(state.completedAt && { completedAt: state.completedAt }),
+    }
+  }
+
+  /**
+   * Wait for a tracked job to reach a terminal status by polling its in-memory
+   * state. Adapters that drive the state machine in a background promise can
+   * call this from `waitForCompletion`.
+   */
+  async waitForCompletion(id: string, pollInterval = 1000): Promise<BatchResult[]> {
+    const state = this.get(id)
+    while (
+      state.status !== 'completed' &&
+      state.status !== 'failed' &&
+      state.status !== 'cancelled'
+    ) {
+      await sleep(pollInterval)
+    }
+    return state.results
+  }
+
+  /** For tests: drop everything. */
+  clear(): void {
+    this.jobs.clear()
+    this.counter = 0
+  }
+}
+
+// ============================================================================
+// JSON-Schema conversion (used by Anthropic + OpenAI adapters)
+// ============================================================================
+
+/** Zod schema definition structure for type introspection. */
+interface ZodDef {
+  typeName?: string
+  type?: unknown
+  shape?: () => Record<string, unknown>
+}
+
+/** Zod schema with `_def` property for introspection. */
+interface ZodSchemaLike {
+  _def?: ZodDef
+}
+
+/**
+ * Minimal Zod -> JSON Schema converter.
+ *
+ * This is the same simplified converter that previously lived (duplicated)
+ * inside `anthropic.ts` and `openai.ts`. Extracted here so both adapters call
+ * the same implementation. For richer conversion use `zod-to-json-schema`.
+ */
+export function zodToJsonSchema(zodSchema: unknown): Record<string, unknown> {
+  const schema = zodSchema as ZodSchemaLike
+
+  if (!schema._def) {
+    return { type: 'object' }
+  }
+
+  switch (schema._def.typeName) {
+    case 'ZodString':
+      return { type: 'string' }
+    case 'ZodNumber':
+      return { type: 'number' }
+    case 'ZodBoolean':
+      return { type: 'boolean' }
+    case 'ZodArray':
+      return { type: 'array', items: zodToJsonSchema(schema._def.type) }
+    case 'ZodObject': {
+      const shape = schema._def.shape?.() ?? {}
+      const properties: Record<string, unknown> = {}
+      for (const [key, value] of Object.entries(shape)) {
+        properties[key] = zodToJsonSchema(value)
+      }
+      return { type: 'object', properties, required: Object.keys(properties) }
+    }
+    default:
+      return { type: 'object' }
+  }
+}
+
+// ============================================================================
+// JSON parsing (used by every adapter that returns raw text)
+// ============================================================================
+
+/**
+ * Try to parse `text` as JSON when it looks like JSON or a schema is expected,
+ * otherwise return the text unchanged. Never throws.
+ */
+export function tryParseJson(text: string | undefined, expectJson = false): unknown {
+  if (!text) return text
+  const trimmed = text.trim()
+  const looksLikeJson = trimmed.startsWith('{') || trimmed.startsWith('[')
+  if (!expectJson && !looksLikeJson) return text
+  try {
+    return JSON.parse(text)
+  } catch {
+    return text
+  }
+}
+
+// ============================================================================
+// Misc
+// ============================================================================
+
+/** Promise-based setTimeout. */
+export function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms))
+}

package/src/batch-map.ts

@@ -27,6 +27,7 @@ import {
   isFlexAvailable,
   type ExecutionTier,
 } from './context.js'
+import { getLogger } from './logger.js'
 import { createBatch, getBatchAdapter, type BatchItem, type BatchResult } from './batch-queue.js'
 import { generateObject, generateText } from './generate.js'
 import type { SimpleSchema } from './schema.js'
@@ -100,11 +101,7 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
   /** Cached resolver promise */
   private _resolver: Promise<T[]> | null = null
 
-  constructor(
-    items: unknown[],
-    operations: CapturedOperation[][],
-    options: BatchMapOptions = {}
-  ) {
+  constructor(items: unknown[], operations: CapturedOperation[][], options: BatchMapOptions = {}) {
     this._items = items
     this._operations = operations
     this._options = options
@@ -144,7 +141,9 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
         if (isFlexAvailable()) {
           return this._resolveViaFlex()
         }
-        console.warn(`Flex processing not available for ${getProvider()}, using immediate execution`)
+        getLogger().warn(
+          `Flex processing not available for ${getProvider()}, using immediate execution`
+        )
         return this._resolveImmediately()
 
       case 'batch':
@@ -186,8 +185,8 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
           batchItems.push({
             id,
             prompt,
-            schema: op.schema,
-            options: { system: op.system, model },
+            ...(op.schema !== undefined && { schema: op.schema }),
+            options: Object.assign({ model }, op.system !== undefined ? { system: op.system } : {}),
             status: 'pending',
           })
 
@@ -200,7 +199,7 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
       return this._reconstructResults(results, itemOperationMap)
     } catch {
       // Flex adapter not available, fall back to batch or immediate
-      console.warn(`Flex adapter not available, falling back to batch API`)
+      getLogger().warn(`Flex adapter not available, falling back to batch API`)
       return this._resolveViaBatchAPI()
     }
   }
@@ -219,7 +218,9 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
       adapter = getBatchAdapter(provider)
     } catch {
       // Adapter not registered, fall back to immediate execution
-      console.warn(`Batch adapter for ${provider} not available, falling back to immediate execution`)
+      getLogger().warn(
+        `Batch adapter for ${provider} not available, falling back to immediate execution`
+      )
       return this._resolveImmediately()
     }
 
@@ -243,8 +244,8 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
         batchItems.push({
           id,
           prompt,
-          schema: op.schema,
-          options: { system: op.system, model },
+          ...(op.schema !== undefined && { schema: op.schema }),
+          options: Object.assign({ model }, op.system !== undefined ? { system: op.system } : {}),
           status: 'pending',
         })
 
@@ -256,15 +257,15 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
     const batch = createBatch({
       provider,
       model,
-      webhookUrl: ctx.webhookUrl,
-      metadata: ctx.metadata,
+      ...(ctx.webhookUrl !== undefined && { webhookUrl: ctx.webhookUrl }),
+      ...(ctx.metadata !== undefined && { metadata: ctx.metadata }),
     })
 
     for (const item of batchItems) {
       batch.add(item.prompt, {
-        schema: item.schema,
-        options: item.options,
-        customId: item.id,
+        ...(item.schema !== undefined && { schema: item.schema }),
+        ...(item.options !== undefined && { options: item.options }),
+        ...(item.id !== undefined && { customId: item.id }),
       })
     }
 
@@ -324,7 +325,11 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
   ): Promise<unknown> {
     switch (op.type) {
       case 'text':
-        const textResult = await generateText({ model, prompt, system: op.system })
+        const textResult = await generateText({
+          model,
+          prompt,
+          ...(op.system !== undefined && { system: op.system }),
+        })
         return textResult.text
 
       case 'boolean':
@@ -341,7 +346,7 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
           model,
           schema: { items: ['List items'] },
           prompt,
-          system: op.system,
+          ...(op.system !== undefined && { system: op.system }),
         })
         return (listResult.object as { items: string[] }).items
 
@@ -351,7 +356,7 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
           model,
           schema: op.schema || { result: 'The result' },
           prompt,
-          system: op.system,
+          ...(op.system !== undefined && { system: op.system }),
         })
         return objResult.object
     }
@@ -387,7 +392,7 @@ export class BatchMapPromise<T> implements PromiseLike<T[]> {
       if (operations.length === 1) {
         results[itemIndex] = batchResult.result as T
       } else {
-        (results[itemIndex] as Record<string, unknown>)[`op_${opIndex}`] = batchResult.result
+        ;(results[itemIndex] as Record<string, unknown>)[`op_${opIndex}`] = batchResult.result
       }
     }
 
@@ -499,9 +504,9 @@ export function captureOperation(
     id: `op_${++operationCounter}`,
     prompt,
     itemPlaceholder: currentItemPlaceholder,
-    schema,
+    ...(schema !== undefined && { schema }),
     type,
-    system,
+    ...(system !== undefined && { system }),
   })
 }
 
@@ -571,6 +576,6 @@ export function isBatchMapPromise(value: unknown): value is BatchMapPromise<unkn
     value !== null &&
     typeof value === 'object' &&
     BATCH_MAP_SYMBOL in value &&
-    (value as any)[BATCH_MAP_SYMBOL] === true
+    (value as Record<symbol, unknown>)[BATCH_MAP_SYMBOL] === true
   )
 }