ai-functions 2.1.3 → 2.4.0

package/src/batch/cloudflare.ts

@@ -1,11 +1,11 @@
 /**
- * Cloudflare AI Gateway Batch Adapter
+ * Cloudflare AI Gateway Adapter
  *
- * Implements batch processing using Cloudflare AI Gateway's batch endpoint.
- * Cloudflare's batch API is newer but integrates well with Workers.
+ * Cloudflare's AI Gateway doesn't have a native batch API like OpenAI/Anthropic,
+ * so this adapter fakes batch processing via concurrent direct calls through
+ * the gateway and tracks state locally (`LocalJobStore` from `./provider.js`).
  *
- * Note: Cloudflare's batch API is evolving. This adapter implements
- * the current batch capabilities and can be extended as new features land.
+ * For true async batch processing, consider Cloudflare Queues + Workers.
  *
  * @see https://developers.cloudflare.com/ai-gateway/
  *
@@ -13,325 +13,149 @@
  */
 
 import {
+  LocalJobStore,
+  processConcurrently,
   registerBatchAdapter,
+  tryParseJson,
   type BatchAdapter,
   type BatchItem,
   type BatchJob,
   type BatchQueueOptions,
   type BatchResult,
   type BatchSubmitResult,
-  type BatchStatus,
-} from '../batch-queue.js'
-import { schema as convertSchema } from '../schema.js'
+} from './provider.js'
 
 // ============================================================================
-// Types
-// ============================================================================
-
-interface CloudflareBatchRequest {
-  id: string
-  provider: string
-  endpoint: string
-  headers: Record<string, string>
-  query: {
-    model: string
-    messages: Array<{ role: string; content: string }>
-    response_format?: { type: 'json_object' }
-    max_tokens?: number
-    temperature?: number
-  }
-}
-
-interface CloudflareBatchResponse {
-  id: string
-  success: boolean
-  result?: {
-    response: string
-    usage?: {
-      prompt_tokens: number
-      completion_tokens: number
-      total_tokens: number
-    }
-  }
-  error?: {
-    message: string
-    code?: string
-  }
-}
-
-interface CloudflareBatchJob {
-  id: string
-  status: 'pending' | 'processing' | 'completed' | 'failed'
-  created_at: string
-  completed_at?: string
-  request_count: number
-  completed_count: number
-  failed_count: number
-  results_url?: string
-}
-
-// ============================================================================
-// Cloudflare Client Configuration
+// Cloudflare client configuration
 // ============================================================================
 
 let accountId: string | undefined
 let gatewayId: string | undefined
 let apiToken: string | undefined
-let baseUrl = 'https://api.cloudflare.com/client/v4'
 
-/**
- * Configure the Cloudflare client
- */
+/** Configure the Cloudflare client. */
 export function configureCloudflare(options: {
   accountId?: string
   gatewayId?: string
   apiToken?: string
+  /** Reserved for tests / custom hosts; the gateway URL itself is fixed. */
   baseUrl?: string
 }): void {
   if (options.accountId) accountId = options.accountId
   if (options.gatewayId) gatewayId = options.gatewayId
   if (options.apiToken) apiToken = options.apiToken
-  if (options.baseUrl) baseUrl = options.baseUrl
+  // baseUrl is accepted for backwards compatibility but unused.
+  void options.baseUrl
+}
+
+interface CloudflareConfig {
+  accountId: string
+  gatewayId: string
+  apiToken: string
 }
 
-function getConfig(): { accountId: string; gatewayId: string; apiToken: string } {
-  const accId = accountId || process.env.CLOUDFLARE_ACCOUNT_ID
-  const gwId = gatewayId || process.env.CLOUDFLARE_AI_GATEWAY_ID || process.env.AI_GATEWAY_ID
-  const token = apiToken || process.env.CLOUDFLARE_API_TOKEN
+function getConfig(): CloudflareConfig {
+  const accId = accountId || process.env['CLOUDFLARE_ACCOUNT_ID']
+  const gwId = gatewayId || process.env['CLOUDFLARE_AI_GATEWAY_ID'] || process.env['AI_GATEWAY_ID']
+  const token = apiToken || process.env['CLOUDFLARE_API_TOKEN']
 
   if (!accId) {
-    throw new Error('Cloudflare account ID not configured. Set CLOUDFLARE_ACCOUNT_ID or call configureCloudflare()')
+    throw new Error(
+      'Cloudflare account ID not configured. Set CLOUDFLARE_ACCOUNT_ID or call configureCloudflare()'
+    )
   }
   if (!gwId) {
-    throw new Error('Cloudflare AI Gateway ID not configured. Set CLOUDFLARE_AI_GATEWAY_ID or call configureCloudflare()')
+    throw new Error(
+      'Cloudflare AI Gateway ID not configured. Set CLOUDFLARE_AI_GATEWAY_ID or call configureCloudflare()'
+    )
   }
   if (!token) {
-    throw new Error('Cloudflare API token not configured. Set CLOUDFLARE_API_TOKEN or call configureCloudflare()')
+    throw new Error(
+      'Cloudflare API token not configured. Set CLOUDFLARE_API_TOKEN or call configureCloudflare()'
+    )
   }
 
   return { accountId: accId, gatewayId: gwId, apiToken: token }
 }
 
-async function cloudflareRequest<T>(
-  method: 'GET' | 'POST',
-  path: string,
-  body?: unknown
-): Promise<T> {
-  const config = getConfig()
-  const url = `${baseUrl}${path}`
-
-  const response = await fetch(url, {
-    method,
-    headers: {
-      Authorization: `Bearer ${config.apiToken}`,
-      'Content-Type': 'application/json',
-    },
-    body: body ? JSON.stringify(body) : undefined,
-  })
-
-  if (!response.ok) {
-    const error = await response.text()
-    throw new Error(`Cloudflare API error: ${response.status} ${error}`)
-  }
-
-  const data = await response.json() as { success: boolean; result: T; errors?: unknown[] }
-  if (!data.success) {
-    throw new Error(`Cloudflare API error: ${JSON.stringify(data.errors)}`)
-  }
-
-  return data.result
-}
-
 // ============================================================================
-// In-memory job storage (for polling)
+// Local job tracking
 // ============================================================================
 
-const pendingJobs = new Map<string, {
-  items: BatchItem[]
-  options: BatchQueueOptions
-  results: BatchResult[]
-  status: BatchStatus
-  createdAt: Date
-  completedAt?: Date
-}>()
-
-let jobCounter = 0
+const jobs = new LocalJobStore('cf_batch')
 
 // ============================================================================
-// Cloudflare Batch Adapter
+// Cloudflare batch adapter (BatchProvider port)
 // ============================================================================
 
-/**
- * Cloudflare batch adapter
- *
- * Note: Cloudflare's AI Gateway doesn't have a native batch API like OpenAI/Anthropic.
- * This adapter implements batch processing by:
- * 1. Sending requests concurrently through the gateway
- * 2. Utilizing Cloudflare's caching and rate limiting
- * 3. Tracking job state locally (or in D1/KV for production)
- *
- * For true async batch processing, consider using Cloudflare Queues + Workers.
- */
 const cloudflareAdapter: BatchAdapter = {
   async submit(items: BatchItem[], options: BatchQueueOptions): Promise<BatchSubmitResult> {
     const config = getConfig()
-    const jobId = `cf_batch_${++jobCounter}_${Date.now()}`
     const model = options.model || 'mistral/mistral-7b-instruct-v0.1'
-
-    // Store job state
-    pendingJobs.set(jobId, {
-      items,
-      options,
-      results: [],
-      status: 'pending',
-      createdAt: new Date(),
-    })
-
-    // Process requests concurrently (Cloudflare handles rate limiting)
-    const completion = processCloudflareRequests(jobId, items, config, model, options)
+    const { id, state } = jobs.create(items, options)
+
+    const completion = (async () => {
+      state.status = 'in_progress'
+      const results = await processConcurrently(
+        items,
+        (item) => processCloudflareItem(item, config, model),
+        {
+          concurrency: 10,
+          onWaveComplete: (partial) => {
+            state.results = partial
+          },
+        }
+      )
+      state.results = results
+      state.status = results.every((r) => r.status === 'completed') ? 'completed' : 'failed'
+      state.completedAt = new Date()
+      return results
+    })()
 
     const job: BatchJob = {
-      id: jobId,
+      id,
       provider: 'cloudflare',
       status: 'pending',
       totalItems: items.length,
       completedItems: 0,
       failedItems: 0,
-      createdAt: new Date(),
-      webhookUrl: options.webhookUrl,
+      createdAt: state.createdAt,
+      ...(options.webhookUrl !== undefined && { webhookUrl: options.webhookUrl }),
     }
 
     return { job, completion }
   },
 
   async getStatus(batchId: string): Promise<BatchJob> {
-    const job = pendingJobs.get(batchId)
-    if (!job) {
-      throw new Error(`Batch not found: ${batchId}`)
-    }
-
-    const completedItems = job.results.filter((r) => r.status === 'completed').length
-    const failedItems = job.results.filter((r) => r.status === 'failed').length
-
-    return {
-      id: batchId,
-      provider: 'cloudflare',
-      status: job.status,
-      totalItems: job.items.length,
-      completedItems,
-      failedItems,
-      createdAt: job.createdAt,
-      completedAt: job.completedAt,
-    }
+    return jobs.snapshot(batchId, 'cloudflare')
   },
 
   async cancel(batchId: string): Promise<void> {
-    const job = pendingJobs.get(batchId)
-    if (job) {
-      job.status = 'cancelled'
+    if (jobs.has(batchId)) {
+      jobs.get(batchId).status = 'cancelled'
     }
   },
 
   async getResults(batchId: string): Promise<BatchResult[]> {
-    const job = pendingJobs.get(batchId)
-    if (!job) {
-      throw new Error(`Batch not found: ${batchId}`)
-    }
-    return job.results
+    return jobs.get(batchId).results
   },
 
   async waitForCompletion(batchId: string, pollInterval = 1000): Promise<BatchResult[]> {
-    const job = pendingJobs.get(batchId)
-    if (!job) {
-      throw new Error(`Batch not found: ${batchId}`)
-    }
-
-    while (job.status !== 'completed' && job.status !== 'failed' && job.status !== 'cancelled') {
-      await new Promise((resolve) => setTimeout(resolve, pollInterval))
-    }
-
-    return job.results
+    return jobs.waitForCompletion(batchId, pollInterval)
   },
 }
 
 // ============================================================================
-// Processing
+// Per-item processing
 // ============================================================================
 
-async function processCloudflareRequests(
-  jobId: string,
-  items: BatchItem[],
-  config: { accountId: string; gatewayId: string; apiToken: string },
-  model: string,
-  options: BatchQueueOptions
-): Promise<BatchResult[]> {
-  const job = pendingJobs.get(jobId)
-  if (!job) {
-    throw new Error(`Job not found: ${jobId}`)
-  }
-
-  job.status = 'in_progress'
-
-  // Process all requests concurrently with concurrency limit
-  const CONCURRENCY = 10
-  const results: BatchResult[] = []
-
-  for (let i = 0; i < items.length; i += CONCURRENCY) {
-    const batch = items.slice(i, i + CONCURRENCY)
-
-    const batchResults = await Promise.all(
-      batch.map(async (item) => {
-        try {
-          const result = await processCloudflareItem(item, config, model)
-          return result
-        } catch (error) {
-          return {
-            id: item.id,
-            customId: item.id,
-            status: 'failed' as const,
-            error: error instanceof Error ? error.message : 'Unknown error',
-          }
-        }
-      })
-    )
-
-    results.push(...batchResults)
-    job.results = results
-  }
-
-  job.status = results.every((r) => r.status === 'completed') ? 'completed' : 'failed'
-  job.completedAt = new Date()
-
-  return results
-}
-
 async function processCloudflareItem(
   item: BatchItem,
-  config: { accountId: string; gatewayId: string; apiToken: string },
+  config: CloudflareConfig,
   model: string
 ): Promise<BatchResult> {
-  // Route through AI Gateway
   const gatewayUrl = `https://gateway.ai.cloudflare.com/v1/${config.accountId}/${config.gatewayId}`
-
-  // Determine provider from model
-  let provider = 'workers-ai'
-  let endpoint = ''
-
-  if (model.startsWith('openai/') || model.startsWith('gpt-')) {
-    provider = 'openai'
-    endpoint = '/chat/completions'
-  } else if (model.startsWith('anthropic/') || model.startsWith('claude-')) {
-    provider = 'anthropic'
-    endpoint = '/messages'
-  } else if (model.startsWith('@cf/') || model.startsWith('workers-ai/')) {
-    provider = 'workers-ai'
-    endpoint = `/ai/run/${model.replace('workers-ai/', '').replace('@cf/', '')}`
-  } else {
-    // Default to OpenAI-compatible
-    provider = 'openai'
-    endpoint = '/chat/completions'
-  }
-
+  const { provider, endpoint } = routeForModel(model)
   const url = `${gatewayUrl}/${provider}${endpoint}`
 
   const messages = [
@@ -343,12 +167,8 @@ async function processCloudflareItem(
     model: model.replace(`${provider}/`, ''),
     messages,
     max_tokens: item.options?.maxTokens || 4096,
-    temperature: item.options?.temperature,
-  }
-
-  // Add JSON mode if schema is provided
-  if (item.schema) {
-    body.response_format = { type: 'json_object' }
+    ...(item.options?.temperature !== undefined && { temperature: item.options.temperature }),
+    ...(item.schema && { response_format: { type: 'json_object' } }),
   }
 
   const response = await fetch(url, {
@@ -365,55 +185,52 @@ async function processCloudflareItem(
     throw new Error(`Cloudflare Gateway error: ${response.status} ${error}`)
   }
 
-  const data = await response.json() as {
+  const data = (await response.json()) as {
     choices?: Array<{ message: { content: string } }>
     content?: Array<{ text: string }>
     response?: string
     usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number }
   }
 
-  // Extract content based on response format
-  let content: string | undefined
-
-  if (data.choices?.[0]?.message?.content) {
-    // OpenAI format
-    content = data.choices[0].message.content
-  } else if (data.content?.[0]?.text) {
-    // Anthropic format
-    content = data.content[0].text
-  } else if (data.response) {
-    // Workers AI format
-    content = data.response
-  }
-
-  let result: unknown = content
-
-  // Try to parse JSON if schema was provided
-  if (item.schema && content) {
-    try {
-      result = JSON.parse(content)
-    } catch {
-      // Keep as string
-    }
-  }
+  // Extract content based on which downstream provider answered.
+  const content =
+    data.choices?.[0]?.message?.content ?? data.content?.[0]?.text ?? data.response ?? undefined
 
   return {
     id: item.id,
     customId: item.id,
     status: 'completed',
-    result,
-    usage: data.usage
-      ? {
-          promptTokens: data.usage.prompt_tokens,
-          completionTokens: data.usage.completion_tokens,
-          totalTokens: data.usage.total_tokens,
-        }
-      : undefined,
+    result: tryParseJson(content, !!item.schema),
+    ...(data.usage && {
+      usage: {
+        promptTokens: data.usage.prompt_tokens,
+        completionTokens: data.usage.completion_tokens,
+        totalTokens: data.usage.total_tokens,
+      },
+    }),
+  }
+}
+
+/** Map a model id to the AI Gateway provider segment + endpoint path. */
+function routeForModel(model: string): { provider: string; endpoint: string } {
+  if (model.startsWith('openai/') || model.startsWith('gpt-')) {
+    return { provider: 'openai', endpoint: '/chat/completions' }
+  }
+  if (model.startsWith('anthropic/') || model.startsWith('claude-')) {
+    return { provider: 'anthropic', endpoint: '/messages' }
+  }
+  if (model.startsWith('@cf/') || model.startsWith('workers-ai/')) {
+    return {
+      provider: 'workers-ai',
+      endpoint: `/ai/run/${model.replace('workers-ai/', '').replace('@cf/', '')}`,
+    }
   }
+  // Default: assume an OpenAI-compatible downstream.
+  return { provider: 'openai', endpoint: '/chat/completions' }
 }
 
 // ============================================================================
-// Register Adapter
+// Register adapter
 // ============================================================================
 
 registerBatchAdapter('cloudflare', cloudflareAdapter)