ai-functions 0.3.0 → 2.0.1

package/test/e2e-flex.test.ts

@@ -0,0 +1,209 @@
+/**
+ * E2E Tests for Flex Tier Processing
+ *
+ * These tests hit real APIs to verify the flex tier actually works.
+ * Requires OPENAI_API_KEY environment variable.
+ *
+ * Run with: npx vitest run test/e2e-flex.test.ts
+ */
+
+import { describe, it, expect, beforeAll, afterAll } from 'vitest'
+import { configure, resetContext, getExecutionTier, isFlexAvailable } from '../src/context.js'
+import { getFlexAdapter, type BatchItem, type BatchResult } from '../src/batch-queue.js'
+
+// Import the OpenAI adapter to register it
+import { configureOpenAI } from '../src/batch/openai.js'
+import '../src/batch/openai.js'
+
+// Skip tests if no API key (either direct or via gateway)
+const hasApiKey = !!(process.env.OPENAI_API_KEY || (process.env.AI_GATEWAY_URL && process.env.AI_GATEWAY_TOKEN))
+
+describe.skipIf(!hasApiKey)('E2E Flex Tier Processing', () => {
+  beforeAll(() => {
+    // If using gateway, configure the base URL
+    if (process.env.AI_GATEWAY_URL && process.env.AI_GATEWAY_TOKEN && !process.env.OPENAI_API_KEY) {
+      configureOpenAI({
+        apiKey: process.env.AI_GATEWAY_TOKEN,
+        baseUrl: `${process.env.AI_GATEWAY_URL}/openai`,
+      })
+    }
+
+    configure({
+      provider: 'openai',
+      model: 'gpt-4o-mini', // Use mini for faster/cheaper tests
+      batchMode: 'auto',
+      flexThreshold: 3,
+      batchThreshold: 500,
+    })
+  })
+
+  afterAll(() => {
+    resetContext()
+  })
+
+  describe('Flex Adapter Direct', () => {
+    it('processes a small batch via flex adapter', async () => {
+      const adapter = getFlexAdapter('openai')
+
+      const items: BatchItem[] = [
+        {
+          id: 'item_1',
+          prompt: 'What is 2 + 2? Reply with just the number.',
+          status: 'pending',
+        },
+        {
+          id: 'item_2',
+          prompt: 'What is 3 + 3? Reply with just the number.',
+          status: 'pending',
+        },
+        {
+          id: 'item_3',
+          prompt: 'What is 4 + 4? Reply with just the number.',
+          status: 'pending',
+        },
+      ]
+
+      const results = await adapter.submitFlex(items, { model: 'gpt-4o-mini' })
+
+      expect(results).toHaveLength(3)
+      results.forEach((result) => {
+        expect(result.status).toBe('completed')
+        expect(result.result).toBeDefined()
+      })
+
+      // Check the actual responses contain numbers
+      const result1 = results.find((r) => r.id === 'item_1')
+      const result2 = results.find((r) => r.id === 'item_2')
+      const result3 = results.find((r) => r.id === 'item_3')
+
+      expect(result1?.result).toContain('4')
+      expect(result2?.result).toContain('6')
+      expect(result3?.result).toContain('8')
+    }, 30000) // 30 second timeout
+
+    it('processes items with structured output (JSON schema)', async () => {
+      const adapter = getFlexAdapter('openai')
+
+      const items: BatchItem[] = [
+        {
+          id: 'structured_1',
+          prompt: 'Generate a person with name "Alice" and age 30',
+          schema: { name: 'string', age: 'number' },
+          status: 'pending',
+        },
+        {
+          id: 'structured_2',
+          prompt: 'Generate a person with name "Bob" and age 25',
+          schema: { name: 'string', age: 'number' },
+          status: 'pending',
+        },
+      ]
+
+      const results = await adapter.submitFlex(items, { model: 'gpt-4o-mini' })
+
+      expect(results).toHaveLength(2)
+      results.forEach((result) => {
+        expect(result.status).toBe('completed')
+        expect(result.result).toBeDefined()
+        expect(typeof result.result).toBe('object')
+      })
+
+      const result1 = results.find((r) => r.id === 'structured_1')
+      const result2 = results.find((r) => r.id === 'structured_2')
+
+      // Results should be parsed JSON objects
+      expect((result1?.result as any)?.name).toBeDefined()
+      expect((result2?.result as any)?.name).toBeDefined()
+    }, 30000)
+
+    it('handles errors gracefully', async () => {
+      const adapter = getFlexAdapter('openai')
+
+      const items: BatchItem[] = [
+        {
+          id: 'valid',
+          prompt: 'Say hello',
+          status: 'pending',
+        },
+      ]
+
+      // This should work even with a minimal prompt
+      const results = await adapter.submitFlex(items, { model: 'gpt-4o-mini' })
+
+      expect(results).toHaveLength(1)
+      expect(results[0].status).toBe('completed')
+    }, 30000)
+
+    it('reports token usage', async () => {
+      const adapter = getFlexAdapter('openai')
+
+      const items: BatchItem[] = [
+        {
+          id: 'usage_test',
+          prompt: 'Count from 1 to 5',
+          status: 'pending',
+        },
+      ]
+
+      const results = await adapter.submitFlex(items, { model: 'gpt-4o-mini' })
+
+      expect(results[0].usage).toBeDefined()
+      expect(results[0].usage?.promptTokens).toBeGreaterThan(0)
+      expect(results[0].usage?.completionTokens).toBeGreaterThan(0)
+      expect(results[0].usage?.totalTokens).toBeGreaterThan(0)
+    }, 30000)
+  })
+
+  describe('Execution Tier Selection', () => {
+    it('selects immediate tier for < flexThreshold items', () => {
+      configure({ batchMode: 'auto', flexThreshold: 5, batchThreshold: 500 })
+      expect(getExecutionTier(1)).toBe('immediate')
+      expect(getExecutionTier(4)).toBe('immediate')
+    })
+
+    it('selects flex tier for flexThreshold to < batchThreshold items', () => {
+      configure({ batchMode: 'auto', flexThreshold: 5, batchThreshold: 500 })
+      expect(getExecutionTier(5)).toBe('flex')
+      expect(getExecutionTier(100)).toBe('flex')
+      expect(getExecutionTier(499)).toBe('flex')
+    })
+
+    it('selects batch tier for >= batchThreshold items', () => {
+      configure({ batchMode: 'auto', flexThreshold: 5, batchThreshold: 500 })
+      expect(getExecutionTier(500)).toBe('batch')
+      expect(getExecutionTier(1000)).toBe('batch')
+    })
+
+    it('reports flex as available for openai', () => {
+      configure({ provider: 'openai' })
+      expect(isFlexAvailable()).toBe(true)
+    })
+  })
+
+  describe('Concurrent Processing', () => {
+    it('processes multiple items concurrently', async () => {
+      const adapter = getFlexAdapter('openai')
+
+      // Create 10 items to test concurrent processing
+      const items: BatchItem[] = Array.from({ length: 10 }, (_, i) => ({
+        id: `concurrent_${i}`,
+        prompt: `What is ${i} + 1? Reply with just the number.`,
+        status: 'pending' as const,
+      }))
+
+      const startTime = Date.now()
+      const results = await adapter.submitFlex(items, { model: 'gpt-4o-mini' })
+      const duration = Date.now() - startTime
+
+      expect(results).toHaveLength(10)
+      results.forEach((result, i) => {
+        expect(result.status).toBe('completed')
+      })
+
+      // With concurrency of 10, this should complete much faster than sequential
+      // Sequential would be ~10-20 seconds, concurrent should be ~2-5 seconds
+      console.log(`Processed 10 items in ${duration}ms`)
+      expect(duration).toBeLessThan(20000) // Should be well under 20 seconds
+    }, 60000)
+  })
+})

package/test/e2e-google-manual.ts

@@ -0,0 +1,178 @@
+#!/usr/bin/env npx tsx
+/**
+ * Manual E2E Test for Google GenAI (Gemini) Flex Tier Processing
+ *
+ * Run with your Google API key:
+ *
+ *   GOOGLE_API_KEY=... npx tsx test/e2e-google-manual.ts
+ *
+ * Or:
+ *
+ *   GEMINI_API_KEY=... npx tsx test/e2e-google-manual.ts
+ */
+
+import { configure, resetContext, getExecutionTier, isFlexAvailable, getProvider } from '../src/context.js'
+import { type BatchItem } from '../src/batch-queue.js'
+import { configureGoogleGenAI, googleFlexAdapter } from '../src/batch/google.js'
+
+// Gemini models
+const GEMINI_FLASH = 'gemini-2.0-flash'
+const GEMINI_PRO = 'gemini-1.5-pro'
+
+async function main() {
+  console.log('\n🧪 E2E Google GenAI (Gemini) Flex Tier Test\n')
+
+  // Check for API key
+  const apiKey = process.env.GOOGLE_API_KEY || process.env.GEMINI_API_KEY
+  if (!apiKey) {
+    console.error('❌ Google API key required')
+    console.error('   Set GOOGLE_API_KEY or GEMINI_API_KEY')
+    console.error('   Get one at: https://aistudio.google.com/app/apikey')
+    process.exit(1)
+  }
+
+  console.log('✅ Google API key found')
+
+  // Configure
+  configure({
+    provider: 'google',
+    model: GEMINI_FLASH,
+    batchMode: 'auto',
+    flexThreshold: 3,
+    batchThreshold: 500,
+  })
+
+  console.log(`✅ Configured: provider=${getProvider()}, flexAvailable=${isFlexAvailable()}`)
+
+  // Test 1: Basic flex processing with Gemini 2.0 Flash
+  console.log('\n📝 Test 1: Basic Flex Processing (Gemini 2.0 Flash)')
+  console.log('   Creating 3 simple prompts...')
+
+  const items: BatchItem[] = [
+    { id: 'test_1', prompt: 'What is 2 + 2? Reply with just the number.', status: 'pending' },
+    { id: 'test_2', prompt: 'What is 3 + 3? Reply with just the number.', status: 'pending' },
+    { id: 'test_3', prompt: 'What is 4 + 4? Reply with just the number.', status: 'pending' },
+  ]
+
+  console.log('   Submitting via flex adapter...')
+  const startTime = Date.now()
+
+  try {
+    const results = await googleFlexAdapter.submitFlex(items, { model: GEMINI_FLASH })
+    const duration = Date.now() - startTime
+
+    console.log(`   ✅ Completed in ${duration}ms`)
+    console.log('   Results:')
+
+    results.forEach((result) => {
+      const status = result.status === 'completed' ? '✅' : '❌'
+      console.log(`     ${status} ${result.id}: "${result.result}"`)
+      if (result.error) {
+        console.log(`        Error: ${result.error}`)
+      }
+    })
+  } catch (error) {
+    console.error(`   ❌ Error: ${error}`)
+    process.exit(1)
+  }
+
+  // Test 2: Token usage
+  console.log('\n📝 Test 2: Token Usage Tracking')
+
+  const usageItems: BatchItem[] = [
+    { id: 'usage', prompt: 'Count from 1 to 5', status: 'pending' },
+  ]
+
+  try {
+    const results = await googleFlexAdapter.submitFlex(usageItems, { model: GEMINI_FLASH })
+    const usage = results[0].usage
+
+    if (usage) {
+      console.log(`   Prompt tokens: ${usage.promptTokens}`)
+      console.log(`   Completion tokens: ${usage.completionTokens}`)
+      console.log(`   Total tokens: ${usage.totalTokens}`)
+      console.log('   ✅ Token usage tracking works')
+    } else {
+      console.log('   ⚠️  No usage data returned')
+    }
+  } catch (error) {
+    console.error(`   ❌ Error: ${error}`)
+  }
+
+  // Test 3: Structured output (JSON)
+  console.log('\n📝 Test 3: Structured Output (JSON)')
+
+  const jsonItems: BatchItem[] = [
+    {
+      id: 'json_test',
+      prompt: 'Generate a JSON object with name "Alice" and age 30. Return only valid JSON, no markdown.',
+      schema: { name: 'string', age: 'number' },
+      status: 'pending',
+    },
+  ]
+
+  try {
+    const results = await googleFlexAdapter.submitFlex(jsonItems, { model: GEMINI_FLASH })
+    console.log(`   Result: ${JSON.stringify(results[0].result)}`)
+    console.log('   ✅ Structured output works')
+  } catch (error) {
+    console.error(`   ❌ Error: ${error}`)
+  }
+
+  // Test 4: Concurrent processing
+  console.log('\n📝 Test 4: Concurrent Processing (10 items)')
+
+  const concurrentItems: BatchItem[] = Array.from({ length: 10 }, (_, i) => ({
+    id: `concurrent_${i}`,
+    prompt: `What is ${i + 1} * 2? Reply with just the number.`,
+    status: 'pending' as const,
+  }))
+
+  try {
+    const startTime = Date.now()
+    const results = await googleFlexAdapter.submitFlex(concurrentItems, { model: GEMINI_FLASH })
+    const duration = Date.now() - startTime
+
+    const successCount = results.filter((r) => r.status === 'completed').length
+    console.log(`   ✅ Processed ${successCount}/10 items in ${duration}ms`)
+    console.log(`   Average: ${Math.round(duration / 10)}ms per item`)
+
+    results.forEach((result) => {
+      console.log(`     ${result.id}: "${result.result}"`)
+    })
+  } catch (error) {
+    console.error(`   ❌ Error: ${error}`)
+  }
+
+  // Test 5: Gemini 1.5 Pro
+  console.log('\n📝 Test 5: Gemini 1.5 Pro')
+
+  const proItems: BatchItem[] = [
+    { id: 'pro_test', prompt: 'Write a haiku about AI.', status: 'pending' },
+  ]
+
+  try {
+    const results = await googleFlexAdapter.submitFlex(proItems, { model: GEMINI_PRO })
+    console.log(`   ✅ Gemini 1.5 Pro response:`)
+    console.log(`   ${results[0].result}`)
+  } catch (error) {
+    console.error(`   ❌ Error: ${error}`)
+  }
+
+  // Test 6: Tier selection
+  console.log('\n📝 Test 6: Execution Tier Selection')
+
+  console.log(`   1 item   → ${getExecutionTier(1)} (expected: immediate)`)
+  console.log(`   3 items  → ${getExecutionTier(3)} (expected: flex)`)
+  console.log(`   10 items → ${getExecutionTier(10)} (expected: flex)`)
+  console.log(`   500 items → ${getExecutionTier(500)} (expected: batch)`)
+
+  console.log('\n✨ E2E Tests Complete!\n')
+
+  resetContext()
+}
+
+main().catch((error) => {
+  console.error('Fatal error:', error)
+  process.exit(1)
+})

package/test/e2e-google.test.ts

@@ -0,0 +1,216 @@
+/**
+ * E2E Tests for Google GenAI (Gemini) Flex Tier Processing
+ *
+ * These tests hit real Google GenAI APIs to verify the flex tier actually works.
+ * Requires Google API key:
+ * - GOOGLE_API_KEY or GEMINI_API_KEY
+ *
+ * Run with: npx vitest run test/e2e-google.test.ts
+ */
+
+import { describe, it, expect, beforeAll, afterAll } from 'vitest'
+import { configure, resetContext, getExecutionTier, isFlexAvailable } from '../src/context.js'
+import { getFlexAdapter, type BatchItem } from '../src/batch-queue.js'
+
+// Import the Google adapter to register it
+import { configureGoogleGenAI, googleFlexAdapter } from '../src/batch/google.js'
+import '../src/batch/google.js'
+
+// Skip tests if no Google API key or gateway
+const hasApiKey = !!(
+  process.env.GOOGLE_API_KEY ||
+  process.env.GEMINI_API_KEY ||
+  (process.env.AI_GATEWAY_URL && process.env.AI_GATEWAY_TOKEN)
+)
+
+// Gemini 2.0 Flash model
+const GEMINI_MODEL = 'gemini-2.0-flash'
+
+describe.skipIf(!hasApiKey)('E2E Google GenAI Flex Tier Processing', () => {
+  beforeAll(() => {
+    // Configure Google GenAI - can use either direct API key or AI Gateway
+    configureGoogleGenAI({
+      gatewayUrl: process.env.AI_GATEWAY_URL,
+      gatewayToken: process.env.AI_GATEWAY_TOKEN,
+    })
+
+    configure({
+      provider: 'google',
+      model: GEMINI_MODEL,
+      batchMode: 'auto',
+      flexThreshold: 3,
+      batchThreshold: 500,
+    })
+  })
+
+  afterAll(() => {
+    resetContext()
+  })
+
+  describe('Flex Adapter Direct', () => {
+    it('processes a small batch via flex adapter', async () => {
+      const items: BatchItem[] = [
+        {
+          id: 'item_1',
+          prompt: 'What is 2 + 2? Reply with just the number.',
+          status: 'pending',
+        },
+        {
+          id: 'item_2',
+          prompt: 'What is 3 + 3? Reply with just the number.',
+          status: 'pending',
+        },
+        {
+          id: 'item_3',
+          prompt: 'What is 4 + 4? Reply with just the number.',
+          status: 'pending',
+        },
+      ]
+
+      const results = await googleFlexAdapter.submitFlex(items, { model: GEMINI_MODEL })
+
+      expect(results).toHaveLength(3)
+      results.forEach((result) => {
+        expect(result.status).toBe('completed')
+        expect(result.result).toBeDefined()
+      })
+
+      // Check the actual responses contain numbers
+      const result1 = results.find((r) => r.id === 'item_1')
+      const result2 = results.find((r) => r.id === 'item_2')
+      const result3 = results.find((r) => r.id === 'item_3')
+
+      expect(String(result1?.result)).toContain('4')
+      expect(String(result2?.result)).toContain('6')
+      expect(String(result3?.result)).toContain('8')
+    }, 30000)
+
+    it('reports token usage', async () => {
+      const items: BatchItem[] = [
+        {
+          id: 'usage_test',
+          prompt: 'Count from 1 to 5',
+          status: 'pending',
+        },
+      ]
+
+      const results = await googleFlexAdapter.submitFlex(items, { model: GEMINI_MODEL })
+
+      expect(results[0].usage).toBeDefined()
+      expect(results[0].usage?.promptTokens).toBeGreaterThan(0)
+      expect(results[0].usage?.completionTokens).toBeGreaterThan(0)
+      expect(results[0].usage?.totalTokens).toBeGreaterThan(0)
+    }, 30000)
+
+    it('handles system prompts', async () => {
+      const items: BatchItem[] = [
+        {
+          id: 'system_test',
+          prompt: 'What should I say?',
+          options: {
+            system: 'You are a helpful assistant that always responds with exactly "Hello World"',
+          },
+          status: 'pending',
+        },
+      ]
+
+      const results = await googleFlexAdapter.submitFlex(items, { model: GEMINI_MODEL })
+
+      expect(results[0].status).toBe('completed')
+      expect(String(results[0].result).toLowerCase()).toContain('hello')
+    }, 30000)
+
+    it('handles structured output (JSON)', async () => {
+      const items: BatchItem[] = [
+        {
+          id: 'json_test',
+          prompt: 'Generate a JSON object with name "Alice" and age 30. Return only valid JSON.',
+          schema: { name: 'string', age: 'number' },
+          status: 'pending',
+        },
+      ]
+
+      const results = await googleFlexAdapter.submitFlex(items, { model: GEMINI_MODEL })
+
+      expect(results[0].status).toBe('completed')
+      const result = results[0].result as { name?: string; age?: number }
+      expect(result).toBeDefined()
+      // The result should be parsed JSON
+      expect(typeof result === 'object').toBe(true)
+    }, 30000)
+  })
+
+  describe('Execution Tier Selection', () => {
+    it('selects immediate tier for < flexThreshold items', () => {
+      configure({ batchMode: 'auto', flexThreshold: 5, batchThreshold: 500 })
+      expect(getExecutionTier(1)).toBe('immediate')
+      expect(getExecutionTier(4)).toBe('immediate')
+    })
+
+    it('selects flex tier for flexThreshold to < batchThreshold items', () => {
+      configure({ batchMode: 'auto', flexThreshold: 5, batchThreshold: 500 })
+      expect(getExecutionTier(5)).toBe('flex')
+      expect(getExecutionTier(100)).toBe('flex')
+    })
+
+    it('reports flex as available for google', () => {
+      configure({ provider: 'google' })
+      expect(isFlexAvailable()).toBe(true)
+    })
+  })
+
+  describe('Concurrent Processing', () => {
+    it('processes multiple items concurrently', async () => {
+      // Create 10 items to test concurrent processing
+      const items: BatchItem[] = Array.from({ length: 10 }, (_, i) => ({
+        id: `concurrent_${i}`,
+        prompt: `What is ${i} + 1? Reply with just the number.`,
+        status: 'pending' as const,
+      }))
+
+      const startTime = Date.now()
+      const results = await googleFlexAdapter.submitFlex(items, { model: GEMINI_MODEL })
+      const duration = Date.now() - startTime
+
+      expect(results).toHaveLength(10)
+      results.forEach((result) => {
+        expect(result.status).toBe('completed')
+      })
+
+      console.log(`Processed 10 items in ${duration}ms`)
+      // With concurrency, should be much faster than sequential
+      expect(duration).toBeLessThan(30000) // Should complete in under 30 seconds
+    }, 60000)
+  })
+})
+
+// Test with different Gemini models (optional - may not be available via gateway)
+describe.skipIf(!hasApiKey)('E2E Google GenAI with Different Models', () => {
+  it.skip('works with gemini-1.5-flash (may not be in gateway)', async () => {
+    const items: BatchItem[] = [
+      {
+        id: 'flash_test',
+        prompt: 'Say "hello" and nothing else.',
+        status: 'pending',
+      },
+    ]
+
+    const results = await googleFlexAdapter.submitFlex(items, { model: 'gemini-1.5-flash' })
+    expect(results[0].status).toBe('completed')
+    expect(String(results[0].result).toLowerCase()).toContain('hello')
+  }, 30000)
+
+  it.skip('works with gemini-1.5-pro (may not be in gateway)', async () => {
+    const items: BatchItem[] = [
+      {
+        id: 'pro_test',
+        prompt: 'Say "hello" and nothing else.',
+        status: 'pending',
+      },
+    ]
+
+    const results = await googleFlexAdapter.submitFlex(items, { model: 'gemini-1.5-pro' })
+    expect(results[0].status).toBe('completed')
+    expect(String(results[0].result).toLowerCase()).toContain('hello')
+  }, 30000)
+})