ai-functions 0.2.19 → 0.4.0

package/test/ai-proxy.test.ts

@@ -0,0 +1,192 @@
+/**
+ * Tests for the ai proxy and AI() schema functions
+ *
+ * These tests use real AI calls via the Cloudflare AI Gateway.
+ */
+
+import { describe, it, expect, beforeEach } from 'vitest'
+import { ai, AI, functions, withTemplate } from '../src/index.js'
+
+// Skip tests if no gateway configured
+const hasGateway = !!process.env.AI_GATEWAY_URL || !!process.env.ANTHROPIC_API_KEY
+
+describe('ai proxy', () => {
+  beforeEach(() => {
+    functions.clear()
+  })
+
+  it('exposes functions registry', () => {
+    expect(ai.functions).toBeDefined()
+    expect(typeof ai.functions.list).toBe('function')
+    expect(typeof ai.functions.get).toBe('function')
+    expect(typeof ai.functions.set).toBe('function')
+    expect(typeof ai.functions.has).toBe('function')
+    expect(typeof ai.functions.clear).toBe('function')
+    expect(typeof ai.functions.delete).toBe('function')
+  })
+
+  it('exposes define helpers', () => {
+    expect(ai.define).toBeDefined()
+    expect(typeof ai.define).toBe('function')
+    expect(typeof ai.define.generative).toBe('function')
+    expect(typeof ai.define.agentic).toBe('function')
+    expect(typeof ai.define.human).toBe('function')
+    expect(typeof ai.define.code).toBe('function')
+  })
+
+  it('exposes defineFunction', () => {
+    expect(typeof ai.defineFunction).toBe('function')
+  })
+})
+
+describe.skipIf(!hasGateway)('ai proxy auto-define', () => {
+  beforeEach(() => {
+    functions.clear()
+  })
+
+  it('auto-defines a function on first call', async () => {
+    expect(functions.has('greetPerson')).toBe(false)
+
+    const result = await (ai as Record<string, (args: unknown) => Promise<unknown>>).greetPerson({
+      name: 'Alice',
+      style: 'friendly',
+    })
+
+    expect(result).toBeDefined()
+    expect(functions.has('greetPerson')).toBe(true)
+  })
+
+  it('uses cached definition on second call', async () => {
+    // First call - defines the function
+    await (ai as Record<string, (args: unknown) => Promise<unknown>>).capitalizeText({
+      text: 'hello',
+    })
+
+    const fn1 = functions.get('capitalizeText')
+    expect(fn1).toBeDefined()
+
+    // Second call - uses cached definition
+    await (ai as Record<string, (args: unknown) => Promise<unknown>>).capitalizeText({
+      text: 'world',
+    })
+
+    const fn2 = functions.get('capitalizeText')
+    expect(fn1).toBe(fn2) // Same cached function
+  })
+})
+
+describe.skipIf(!hasGateway)('AI() schema functions', () => {
+  it('creates schema-based functions', async () => {
+    const client = AI({
+      sentiment: {
+        sentiment: 'positive | negative | neutral',
+        score: 'Confidence score 0-1 (number)',
+        explanation: 'Brief explanation',
+      },
+    })
+
+    expect(client.sentiment).toBeDefined()
+    expect(typeof client.sentiment).toBe('function')
+  })
+
+  it('generates structured output from schema', async () => {
+    const client = AI({
+      person: {
+        name: 'Full name',
+        age: 'Age (number)',
+        occupation: 'Job title',
+      },
+    })
+
+    const result = await client.person('A software engineer named Alice who is 30')
+
+    expect(result).toBeDefined()
+    expect(typeof result.name).toBe('string')
+    expect(typeof result.age).toBe('number')
+    expect(typeof result.occupation).toBe('string')
+  })
+
+  it('generates nested objects', async () => {
+    const client = AI({
+      profile: {
+        user: {
+          name: 'Name',
+          email: 'Email address',
+        },
+        preferences: {
+          theme: 'light | dark',
+          notifications: 'Enabled? (boolean)',
+        },
+      },
+    })
+
+    const result = await client.profile('User Alice who prefers dark mode and has notifications on')
+
+    expect(result).toBeDefined()
+    expect(result.user).toBeDefined()
+    expect(result.preferences).toBeDefined()
+    expect(['light', 'dark']).toContain(result.preferences.theme)
+    expect(typeof result.preferences.notifications).toBe('boolean')
+  })
+
+  it('generates arrays', async () => {
+    const client = AI({
+      todoList: {
+        title: 'List title',
+        items: ['Todo items'],
+        priority: 'high | medium | low',
+      },
+    })
+
+    const result = await client.todoList('A high priority shopping list with 3 items')
+
+    expect(result).toBeDefined()
+    expect(typeof result.title).toBe('string')
+    expect(Array.isArray(result.items)).toBe(true)
+    expect(result.items.length).toBeGreaterThan(0)
+    expect(['high', 'medium', 'low']).toContain(result.priority)
+  })
+})
+
+describe('withTemplate helper', () => {
+  it('handles regular function calls', () => {
+    const fn = withTemplate((prompt: string) => prompt.toUpperCase())
+
+    const result = fn('hello world')
+    expect(result).toBe('HELLO WORLD')
+  })
+
+  it('handles tagged template literals', () => {
+    const fn = withTemplate((prompt: string) => prompt.toUpperCase())
+
+    const result = fn`hello world`
+    expect(result).toBe('HELLO WORLD')
+  })
+
+  it('handles tagged template literals with interpolation', () => {
+    const fn = withTemplate((prompt: string) => prompt.toUpperCase())
+
+    const name = 'Alice'
+    const result = fn`hello ${name}!`
+    expect(result).toBe('HELLO ALICE!')
+  })
+
+  it('handles multiple interpolations', () => {
+    const fn = withTemplate((prompt: string) => prompt)
+
+    const a = 'one'
+    const b = 'two'
+    const c = 'three'
+    const result = fn`${a}, ${b}, ${c}`
+    expect(result).toBe('one, two, three')
+  })
+
+  it('works with async functions', async () => {
+    const fn = withTemplate(async (prompt: string) => {
+      return `Result: ${prompt}`
+    })
+
+    const result = await fn`async test`
+    expect(result).toBe('Result: async test')
+  })
+})

package/test/async-iterators.test.ts

@@ -0,0 +1,327 @@
+/**
+ * Tests for async iterator support on list and extract
+ *
+ * Functions that return lists can be streamed with `for await`:
+ * - list`...` - streams items as they're generated
+ * - extract`...` - streams extracted items
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest'
+
+// ============================================================================
+// Mock async generators
+// ============================================================================
+
+const mockStreamItems = vi.fn()
+
+/**
+ * Create an async iterable from an array (simulates streaming)
+ */
+async function* createAsyncIterable<T>(items: T[], delayMs = 10): AsyncIterable<T> {
+  for (const item of items) {
+    await new Promise(resolve => setTimeout(resolve, delayMs))
+    yield item
+  }
+}
+
+/**
+ * Mock list function that returns both a promise and an async iterable
+ */
+function createMockStreamingList() {
+  return function list(promptOrStrings: string | TemplateStringsArray, ...args: unknown[]) {
+    let prompt: string
+
+    if (Array.isArray(promptOrStrings) && 'raw' in promptOrStrings) {
+      prompt = (promptOrStrings as TemplateStringsArray).reduce((acc, str, i) => {
+        return acc + str + (args[i] ?? '')
+      }, '')
+    } else {
+      prompt = promptOrStrings as string
+    }
+
+    const items = mockStreamItems(prompt)
+
+    // Return an object that is both a Promise and AsyncIterable
+    const asyncIterable = createAsyncIterable(items)
+
+    const result = {
+      // Promise interface - resolve to full array
+      then: (resolve: (value: string[]) => void, reject?: (error: Error) => void) => {
+        return Promise.resolve(items).then(resolve, reject)
+      },
+
+      // AsyncIterable interface - stream items
+      [Symbol.asyncIterator]: () => asyncIterable[Symbol.asyncIterator](),
+    }
+
+    return result as Promise<string[]> & AsyncIterable<string>
+  }
+}
+
+/**
+ * Mock extract function with streaming support
+ */
+function createMockStreamingExtract() {
+  return function extract(promptOrStrings: string | TemplateStringsArray, ...args: unknown[]) {
+    let prompt: string
+
+    if (Array.isArray(promptOrStrings) && 'raw' in promptOrStrings) {
+      prompt = (promptOrStrings as TemplateStringsArray).reduce((acc, str, i) => {
+        return acc + str + (args[i] ?? '')
+      }, '')
+    } else {
+      prompt = promptOrStrings as string
+    }
+
+    const items = mockStreamItems(prompt)
+    const asyncIterable = createAsyncIterable(items)
+
+    const result = {
+      then: (resolve: (value: string[]) => void, reject?: (error: Error) => void) => {
+        return Promise.resolve(items).then(resolve, reject)
+      },
+      [Symbol.asyncIterator]: () => asyncIterable[Symbol.asyncIterator](),
+    }
+
+    return result as Promise<string[]> & AsyncIterable<string>
+  }
+}
+
+// ============================================================================
+// list() async iterator tests
+// ============================================================================
+
+describe('list() async iteration', () => {
+  beforeEach(() => {
+    mockStreamItems.mockReset()
+  })
+
+  it('can be awaited to get full array', async () => {
+    mockStreamItems.mockReturnValue(['Idea 1', 'Idea 2', 'Idea 3'])
+    const list = createMockStreamingList()
+
+    const result = await list`startup ideas`
+
+    expect(Array.isArray(result)).toBe(true)
+    expect(result).toHaveLength(3)
+  })
+
+  it('can be iterated with for await', async () => {
+    mockStreamItems.mockReturnValue(['Idea 1', 'Idea 2', 'Idea 3'])
+    const list = createMockStreamingList()
+
+    const collected: string[] = []
+    for await (const item of list`startup ideas`) {
+      collected.push(item)
+    }
+
+    expect(collected).toHaveLength(3)
+    expect(collected).toEqual(['Idea 1', 'Idea 2', 'Idea 3'])
+  })
+
+  it('allows early termination with break', async () => {
+    mockStreamItems.mockReturnValue(['Idea 1', 'Billion Dollar Idea', 'Idea 3', 'Idea 4', 'Idea 5'])
+    const list = createMockStreamingList()
+
+    const collected: string[] = []
+    for await (const idea of list`startup ideas`) {
+      collected.push(idea)
+      if (idea.includes('Billion')) {
+        break
+      }
+    }
+
+    // Should have stopped after finding the billion dollar idea
+    expect(collected).toHaveLength(2)
+    expect(collected[1]).toBe('Billion Dollar Idea')
+  })
+
+  it('supports nested iteration pattern from README', async () => {
+    const marketList = createMockStreamingList()
+    const icpList = createMockStreamingList()
+
+    // Simulate the nested pattern
+    mockStreamItems
+      .mockReturnValueOnce(['Market A', 'Market B'])
+      .mockReturnValueOnce(['ICP 1', 'ICP 2'])
+      .mockReturnValueOnce(['ICP 3', 'ICP 4'])
+
+    const results: Array<{ market: string; icp: string }> = []
+
+    for await (const market of marketList`market segments`) {
+      for await (const icp of icpList`customer profiles for ${market}`) {
+        results.push({ market, icp })
+      }
+    }
+
+    expect(results).toHaveLength(4)
+    expect(results[0]).toEqual({ market: 'Market A', icp: 'ICP 1' })
+    expect(results[3]).toEqual({ market: 'Market B', icp: 'ICP 4' })
+  })
+
+  it('processes items as they stream in', async () => {
+    mockStreamItems.mockReturnValue(['Item 1', 'Item 2', 'Item 3'])
+    const list = createMockStreamingList()
+
+    const processedAt: number[] = []
+    const startTime = Date.now()
+
+    for await (const _item of list`items`) {
+      processedAt.push(Date.now() - startTime)
+    }
+
+    // Items should be processed incrementally, not all at once
+    expect(processedAt[1]).toBeGreaterThan(processedAt[0])
+    expect(processedAt[2]).toBeGreaterThan(processedAt[1])
+  })
+})
+
+// ============================================================================
+// extract() async iterator tests
+// ============================================================================
+
+describe('extract() async iteration', () => {
+  beforeEach(() => {
+    mockStreamItems.mockReset()
+  })
+
+  it('can be awaited to get full array', async () => {
+    mockStreamItems.mockReturnValue(['John Smith', 'Jane Doe', 'Bob Wilson'])
+    const extract = createMockStreamingExtract()
+
+    const result = await extract`names from article`
+
+    expect(Array.isArray(result)).toBe(true)
+    expect(result).toHaveLength(3)
+  })
+
+  it('can be iterated with for await', async () => {
+    mockStreamItems.mockReturnValue(['email1@example.com', 'email2@example.com'])
+    const extract = createMockStreamingExtract()
+
+    const collected: string[] = []
+    for await (const email of extract`email addresses from document`) {
+      collected.push(email)
+    }
+
+    expect(collected).toHaveLength(2)
+  })
+
+  it('supports the research + extract pattern from README', async () => {
+    mockStreamItems.mockReturnValue(['Competitor A', 'Competitor B', 'Competitor C'])
+    const extract = createMockStreamingExtract()
+
+    const competitors: string[] = []
+    const marketResearch = 'Report mentioning Competitor A, Competitor B, and Competitor C...'
+
+    for await (const competitor of extract`company names from ${marketResearch}`) {
+      competitors.push(competitor)
+      // In real code, you would do: await research`${competitor} vs ${ourProduct}`
+    }
+
+    expect(competitors).toHaveLength(3)
+  })
+
+  it('allows processing each extraction as it completes', async () => {
+    mockStreamItems.mockReturnValue(['email1@test.com', 'email2@test.com'])
+    const extract = createMockStreamingExtract()
+
+    const notifications: string[] = []
+    for await (const email of extract`emails from document`) {
+      // Simulate sending notification
+      notifications.push(`Notified: ${email}`)
+    }
+
+    expect(notifications).toHaveLength(2)
+    expect(notifications[0]).toBe('Notified: email1@test.com')
+  })
+})
+
+// ============================================================================
+// Combined patterns
+// ============================================================================
+
+describe('combined async iteration patterns', () => {
+  beforeEach(() => {
+    mockStreamItems.mockReset()
+  })
+
+  it('supports the full market research pattern from README', async () => {
+    const list = createMockStreamingList()
+    const extract = createMockStreamingExtract()
+
+    // Mock different results for each call
+    mockStreamItems
+      .mockReturnValueOnce(['Market A']) // markets
+      .mockReturnValueOnce(['ICP 1']) // ICPs for Market A
+      .mockReturnValueOnce(['Blog 1', 'Blog 2']) // blog titles
+
+    const results: string[] = []
+
+    // Simplified version of README example
+    for await (const market of list`market segments for idea`) {
+      for await (const icp of list`customer profiles for ${market}`) {
+        for await (const title of list`blog titles for ${icp}`) {
+          results.push(`${market} > ${icp} > ${title}`)
+        }
+      }
+    }
+
+    expect(results).toHaveLength(2)
+    expect(results[0]).toBe('Market A > ICP 1 > Blog 1')
+    expect(results[1]).toBe('Market A > ICP 1 > Blog 2')
+  })
+})
+
+// ============================================================================
+// Type safety
+// ============================================================================
+
+describe('async iterator type safety', () => {
+  it('list returns string items by default', async () => {
+    mockStreamItems.mockReturnValue(['a', 'b', 'c'])
+    const list = createMockStreamingList()
+
+    for await (const item of list`items`) {
+      expect(typeof item).toBe('string')
+    }
+  })
+
+  it('extract can return typed objects with schema', async () => {
+    // This tests the concept - actual implementation would use schema
+    const items = [
+      { name: 'Acme', role: 'competitor' },
+      { name: 'Beta', role: 'partner' },
+    ]
+
+    mockStreamItems.mockReturnValue(items)
+    const extract = createMockStreamingExtract()
+
+    const collected: Array<{ name: string; role: string }> = []
+    for await (const company of extract`companies from text`) {
+      collected.push(company as { name: string; role: string })
+    }
+
+    expect(collected[0]).toHaveProperty('name')
+    expect(collected[0]).toHaveProperty('role')
+  })
+})
+
+// ============================================================================
+// Error handling
+// ============================================================================
+
+describe('async iterator error handling', () => {
+  it('propagates errors during iteration', async () => {
+    mockStreamItems.mockImplementation(() => {
+      throw new Error('Generation failed')
+    })
+    const list = createMockStreamingList()
+
+    await expect(async () => {
+      for await (const _item of list`items`) {
+        // Should not reach here
+      }
+    }).rejects.toThrow('Generation failed')
+  })
+})