ai-functions 2.1.1 → 2.3.0

package/test/generate-core.test.ts

@@ -2,295 +2,206 @@
  * Tests for the core generate() primitive
  *
  * generate(type, prompt, opts?) is the foundation that all other functions use.
+ * Tests require actual AI calls via the Cloudflare AI Gateway.
  */
 
-import { describe, it, expect, vi, beforeEach } from 'vitest'
+import { describe, it, expect } from 'vitest'
+import { generateObject, generateText } from '../src/generate.js'
+import { z } from 'zod'
 
-// ============================================================================
-// Mock for underlying AI calls
-// ============================================================================
-
-const mockAICall = vi.fn()
-
-// Mock generate implementation
-async function generate(
-  type: string,
-  prompt: string,
-  options?: Record<string, unknown>
-): Promise<unknown> {
-  return mockAICall(type, prompt, options)
-}
+// Skip tests if no gateway configured
+const hasGateway = !!process.env.AI_GATEWAY_URL
 
 // ============================================================================
 // generate(type, prompt, opts) signature tests
 // ============================================================================
 
-describe('generate(type, prompt, opts)', () => {
-  beforeEach(() => {
-    mockAICall.mockReset()
-  })
-
+describe.skipIf(!hasGateway)('generate(type, prompt, opts)', () => {
   describe('type: json', () => {
-    it('generates JSON without schema (AI infers structure)', async () => {
-      mockAICall.mockResolvedValue({
-        competitors: ['Competitor A', 'Competitor B'],
-        marketSize: 1000000,
-        trends: ['AI adoption', 'Cloud migration'],
+    it('generates JSON without explicit schema (AI infers structure)', async () => {
+      const result = await generateObject({
+        model: 'haiku',
+        schema: z.object({
+          competitors: z.array(z.string()).describe('List of competitors'),
+          marketSize: z.number().describe('Estimated market size'),
+        }),
+        prompt:
+          'Provide a simple competitive analysis of the cloud computing market. List 2 competitors and an estimated market size in billions.',
       })
 
-      const result = await generate('json', 'competitive analysis of Acme Corp')
-
-      expect(mockAICall).toHaveBeenCalledWith('json', 'competitive analysis of Acme Corp', undefined)
-      expect(result).toHaveProperty('competitors')
-      expect(result).toHaveProperty('marketSize')
+      expect(result.object).toHaveProperty('competitors')
+      expect(result.object).toHaveProperty('marketSize')
+      expect(Array.isArray(result.object.competitors)).toBe(true)
+      expect(typeof result.object.marketSize).toBe('number')
     })
 
     it('generates JSON with schema (typed, validated)', async () => {
-      mockAICall.mockResolvedValue({
-        name: 'Spaghetti Carbonara',
-        servings: 4,
-        ingredients: ['pasta', 'eggs', 'bacon'],
-        steps: ['Boil pasta', 'Cook bacon', 'Mix eggs'],
+      const result = await generateObject({
+        model: 'haiku',
+        schema: z.object({
+          name: z.string().describe('Recipe name'),
+          servings: z.number().describe('Number of servings'),
+          ingredients: z.array(z.string()).describe('List of ingredients'),
+          steps: z.array(z.string()).describe('Cooking steps'),
+        }),
+        prompt: 'Generate a simple 3-ingredient recipe with 2 steps.',
       })
 
-      const result = await generate('json', 'Italian pasta recipe', {
-        schema: {
-          name: 'Recipe name',
-          servings: 'Number of servings (number)',
-          ingredients: ['List of ingredients'],
-          steps: ['Cooking steps'],
-        },
-      })
-
-      expect(mockAICall).toHaveBeenCalledWith(
-        'json',
-        'Italian pasta recipe',
-        expect.objectContaining({ schema: expect.any(Object) })
-      )
-      expect(result).toHaveProperty('name')
-      expect(result).toHaveProperty('servings')
-      expect(typeof (result as { servings: number }).servings).toBe('number')
+      expect(result.object).toHaveProperty('name')
+      expect(result.object).toHaveProperty('servings')
+      expect(typeof result.object.name).toBe('string')
+      expect(typeof result.object.servings).toBe('number')
+      expect(Array.isArray(result.object.ingredients)).toBe(true)
+      expect(Array.isArray(result.object.steps)).toBe(true)
     })
   })
 
   describe('type: text', () => {
     it('generates plain text', async () => {
-      mockAICall.mockResolvedValue('This is the generated text content.')
-
-      const result = await generate('text', 'Write a paragraph about AI')
+      const result = await generateText({
+        model: 'haiku',
+        prompt: 'Write one sentence about AI.',
+      })
 
-      expect(mockAICall).toHaveBeenCalledWith('text', 'Write a paragraph about AI', undefined)
-      expect(typeof result).toBe('string')
+      expect(typeof result.text).toBe('string')
+      expect(result.text.length).toBeGreaterThan(10)
     })
   })
 
   describe('type: code', () => {
-    it('generates code with language option', async () => {
-      mockAICall.mockResolvedValue(`
-function validateEmail(email: string): boolean {
-  return /^[^@]+@[^@]+\\.[^@]+$/.test(email);
-}
-      `.trim())
-
-      const result = await generate('code', 'email validation function', {
-        language: 'typescript',
+    it('generates code with language specified in prompt', async () => {
+      const result = await generateText({
+        model: 'haiku',
+        system:
+          'You are a code generator. Output only valid TypeScript code, no explanations or markdown.',
+        prompt:
+          'Write a TypeScript function called validateEmail that takes a string and returns boolean.',
       })
 
-      expect(mockAICall).toHaveBeenCalledWith(
-        'code',
-        'email validation function',
-        expect.objectContaining({ language: 'typescript' })
-      )
-      expect(typeof result).toBe('string')
-      expect(result).toContain('function')
+      expect(typeof result.text).toBe('string')
+      expect(result.text).toContain('function')
+      expect(result.text).toMatch(/validateEmail|email/i)
     })
 
     it('generates code in different languages', async () => {
-      mockAICall.mockResolvedValue('def validate_email(email):\n    return "@" in email')
-
-      await generate('code', 'email validation', { language: 'python' })
+      const result = await generateText({
+        model: 'haiku',
+        system:
+          'You are a code generator. Output only valid Python code, no explanations or markdown.',
+        prompt:
+          'Write a Python function called validate_email that takes a string and returns a boolean.',
+      })
 
-      expect(mockAICall).toHaveBeenCalledWith(
-        'code',
-        'email validation',
-        expect.objectContaining({ language: 'python' })
-      )
+      expect(typeof result.text).toBe('string')
+      expect(result.text).toContain('def')
     })
   })
 
   describe('type: markdown', () => {
     it('generates markdown content', async () => {
-      mockAICall.mockResolvedValue('# README\n\n## Features\n\n- Feature 1\n- Feature 2')
-
-      const result = await generate('markdown', 'README for a TypeScript library')
+      const result = await generateText({
+        model: 'haiku',
+        system: 'You write in markdown format.',
+        prompt: 'Write a very short README with a heading and 2 bullet points.',
+      })
 
-      expect(mockAICall).toHaveBeenCalledWith('markdown', 'README for a TypeScript library', undefined)
-      expect(typeof result).toBe('string')
-      expect(result).toContain('#')
+      expect(typeof result.text).toBe('string')
+      expect(result.text).toContain('#')
     })
   })
 
   describe('type: yaml', () => {
     it('generates YAML content', async () => {
-      mockAICall.mockResolvedValue(`
-apiVersion: apps/v1
-kind: Deployment
-metadata:
-  name: my-app
-      `.trim())
-
-      const result = await generate('yaml', 'kubernetes deployment for my-app')
-
-      expect(mockAICall).toHaveBeenCalledWith('yaml', 'kubernetes deployment for my-app', undefined)
-      expect(typeof result).toBe('string')
-      expect(result).toContain('apiVersion')
+      const result = await generateText({
+        model: 'haiku',
+        system: 'You output only valid YAML, no explanations or markdown fences.',
+        prompt: 'Generate a simple YAML config with name: "test-app" and port: 3000.',
+      })
+
+      expect(typeof result.text).toBe('string')
+      expect(result.text.toLowerCase()).toContain('name')
     })
   })
 
   describe('type: list', () => {
     it('generates a list of items', async () => {
-      mockAICall.mockResolvedValue(['Item 1', 'Item 2', 'Item 3'])
-
-      const result = await generate('list', '5 startup ideas')
+      const result = await generateObject({
+        model: 'haiku',
+        schema: z.object({
+          items: z.array(z.string()).describe('List of startup ideas'),
+        }),
+        prompt: 'List exactly 3 startup ideas.',
+      })
 
-      expect(mockAICall).toHaveBeenCalledWith('list', '5 startup ideas', undefined)
-      expect(Array.isArray(result)).toBe(true)
+      expect(Array.isArray(result.object.items)).toBe(true)
+      expect(result.object.items.length).toBe(3)
     })
   })
 
   describe('type: diagram', () => {
     it('generates diagram code', async () => {
-      mockAICall.mockResolvedValue('graph TD\n  A[Start] --> B[Process]\n  B --> C[End]')
-
-      const result = await generate('diagram', 'user flow for authentication', {
-        format: 'mermaid',
-      })
-
-      expect(mockAICall).toHaveBeenCalledWith(
-        'diagram',
-        'user flow for authentication',
-        expect.objectContaining({ format: 'mermaid' })
-      )
-      expect(typeof result).toBe('string')
-    })
-  })
-
-  describe('type: slides', () => {
-    it('generates presentation slides', async () => {
-      mockAICall.mockResolvedValue('---\ntheme: default\n---\n\n# Title\n\nContent')
-
-      const result = await generate('slides', 'quarterly review presentation', {
-        format: 'slidev',
-        slides: 10,
+      const result = await generateText({
+        model: 'haiku',
+        system:
+          'You generate Mermaid diagram code. Output only the diagram code, no explanations or markdown fences.',
+        prompt: 'Create a simple flowchart: Start -> Login -> Dashboard.',
       })
 
-      expect(mockAICall).toHaveBeenCalledWith(
-        'slides',
-        'quarterly review presentation',
-        expect.objectContaining({ format: 'slidev', slides: 10 })
-      )
+      expect(typeof result.text).toBe('string')
+      // Mermaid diagrams typically contain --> or -> for connections
+      expect(result.text).toMatch(/(flowchart|graph|-->|->)/i)
     })
   })
 })
 
-// ============================================================================
-// Tagged template support on generate
-// ============================================================================
-
-describe('generate as tagged template', () => {
-  beforeEach(() => {
-    mockAICall.mockReset()
-  })
-
-  // Note: This tests the concept - actual implementation would need the template wrapper
-  it('should support tagged template syntax (conceptual)', async () => {
-    mockAICall.mockResolvedValue({ analysis: 'Result' })
-
-    // This would be: generate`analysis of ${company}`
-    const company = 'Acme Corp'
-    const prompt = `analysis of ${company}`
-
-    const result = await generate('json', prompt)
-
-    expect(result).toHaveProperty('analysis')
-  })
-
-  it('should convert objects to YAML in templates (conceptual)', async () => {
-    mockAICall.mockResolvedValue('Generated content')
-
-    const context = {
-      brand: 'TechCo',
-      audience: 'developers',
-    }
-
-    // This would be: generate`content for ${{ context }}`
-    // The object would be converted to YAML
-    const prompt = `content for\ncontext:\n  brand: TechCo\n  audience: developers`
-
-    await generate('text', prompt)
-
-    expect(mockAICall).toHaveBeenCalledWith('text', expect.stringContaining('brand: TechCo'), undefined)
-  })
-})
-
 // ============================================================================
 // Options parameter tests
 // ============================================================================
 
-describe('generate options', () => {
-  beforeEach(() => {
-    mockAICall.mockReset()
-    mockAICall.mockResolvedValue('result')
-  })
-
-  it('passes model option', async () => {
-    await generate('text', 'test', { model: 'claude-opus-4-5' })
-
-    expect(mockAICall).toHaveBeenCalledWith(
-      'text',
-      'test',
-      expect.objectContaining({ model: 'claude-opus-4-5' })
-    )
-  })
-
-  it('passes temperature option', async () => {
-    await generate('text', 'creative writing', { temperature: 0.9 })
-
-    expect(mockAICall).toHaveBeenCalledWith(
-      'text',
-      'creative writing',
-      expect.objectContaining({ temperature: 0.9 })
-    )
-  })
+describe.skipIf(!hasGateway)('generate options', () => {
+  it('respects temperature option (low temperature = more deterministic)', async () => {
+    // Low temperature should give consistent results
+    const result1 = await generateText({
+      model: 'haiku',
+      prompt: 'Say exactly "hello" and nothing else.',
+      temperature: 0,
+    })
 
-  it('passes maxTokens option', async () => {
-    await generate('text', 'long article', { maxTokens: 4000 })
+    const result2 = await generateText({
+      model: 'haiku',
+      prompt: 'Say exactly "hello" and nothing else.',
+      temperature: 0,
+    })
 
-    expect(mockAICall).toHaveBeenCalledWith(
-      'text',
-      'long article',
-      expect.objectContaining({ maxTokens: 4000 })
-    )
+    // With temperature 0, responses should be very similar
+    expect(result1.text.toLowerCase()).toContain('hello')
+    expect(result2.text.toLowerCase()).toContain('hello')
   })
 
-  it('passes thinking option', async () => {
-    await generate('json', 'complex analysis', { thinking: 'high' })
+  it('accepts maxTokens option without error', async () => {
+    // This test verifies the maxTokens option is passed through without error
+    // The actual truncation behavior is provider-dependent
+    const result = await generateText({
+      model: 'haiku',
+      prompt: 'Say "hello" and nothing else.',
+      maxTokens: 50,
+    })
 
-    expect(mockAICall).toHaveBeenCalledWith(
-      'json',
-      'complex analysis',
-      expect.objectContaining({ thinking: 'high' })
-    )
+    // Just verify we got a response - maxTokens behavior varies by provider/gateway
+    expect(result.text).toBeDefined()
+    expect(typeof result.text).toBe('string')
   })
 
-  it('passes thinking as number (token budget)', async () => {
-    await generate('json', 'complex analysis', { thinking: 10000 })
+  it('passes system prompt correctly', async () => {
+    const result = await generateText({
+      model: 'haiku',
+      system: 'You always respond with exactly one word.',
+      prompt: 'What is your favorite color?',
+    })
 
-    expect(mockAICall).toHaveBeenCalledWith(
-      'json',
-      'complex analysis',
-      expect.objectContaining({ thinking: 10000 })
-    )
+    // With the system prompt, response should be short (ideally one word)
+    const wordCount = result.text.trim().split(/\s+/).length
+    expect(wordCount).toBeLessThanOrEqual(3) // Allow some flexibility
   })
 })
 
@@ -298,22 +209,22 @@ describe('generate options', () => {
 // All convenience functions use generate
 // ============================================================================
 
-describe('convenience functions map to generate', () => {
+describe('convenience functions documentation', () => {
   it('documents the mapping', () => {
     // This test documents the expected mappings
     const mappings = {
-      'ai(prompt)': "generate('text', prompt)",
-      'write(prompt)': "generate('text', prompt)",
-      'code(prompt)': "generate('code', prompt)",
-      'list(prompt)': "generate('list', prompt)",
-      'lists(prompt)': "generate('lists', prompt)",
-      'extract(prompt)': "generate('extract', prompt)",
-      'summarize(prompt)': "generate('summary', prompt)",
-      'diagram(prompt)': "generate('diagram', prompt)",
-      'slides(prompt)': "generate('slides', prompt)",
-      'is(prompt)': "generate('boolean', prompt)",
+      'ai(prompt)': 'generateText({ model, prompt })',
+      'write(prompt)': 'generateText({ model, prompt })',
+      'code(prompt)': "generateText({ model, system: 'code generator', prompt })",
+      'list(prompt)': 'generateObject({ model, schema: { items: [...] }, prompt })',
+      'lists(prompt)':
+        'generateObject({ model, schema: { listName1: [...], listName2: [...] }, prompt })',
+      'extract(prompt)': 'generateObject({ model, schema: { extracted: [...] }, prompt })',
+      'summarize(prompt)': "generateText({ model, system: 'summarizer', prompt })",
+      'diagram(prompt)': "generateText({ model, system: 'mermaid generator', prompt })",
+      'is(prompt)': 'generateObject({ model, schema: { result: boolean }, prompt })',
     }
 
-    expect(Object.keys(mappings)).toHaveLength(10)
+    expect(Object.keys(mappings)).toHaveLength(9)
   })
 })

package/test/implicit-batch.test.ts

@@ -16,7 +16,7 @@
  * ```
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
+import { describe, it, expect, beforeEach, afterEach } from 'vitest'
 import {
   configure,
   resetContext,
@@ -29,54 +29,14 @@ import {
   getFlexThreshold,
   getBatchThreshold,
   isFlexAvailable,
-} from '../src/context.js'
-import { list, write, ai, is } from '../src/primitives.js'
-import {
   createBatchMap,
   BatchMapPromise,
-  captureOperation,
-  isInRecordingMode,
-} from '../src/batch-map.js'
+} from '../src/index.js'
+import { captureOperation } from '../src/batch-map.js'
 
-// Import memory adapter to register it
-import '../src/batch/memory.js'
-import { configureMemoryAdapter, clearBatches } from '../src/batch/memory.js'
-
-// ============================================================================
-// Mock Setup
-// ============================================================================
-
-vi.mock('../src/generate.js', () => ({
-  generateObject: vi.fn().mockImplementation(async ({ prompt, schema }) => {
-    // Simulate list generation
-    if (schema?.items) {
-      return {
-        object: {
-          items: [
-            'Building AI-First Startups in 2026',
-            'The Future of Remote Work',
-            'Sustainable Tech Growth',
-            'From Idea to MVP in 30 Days',
-            'Community-Led Product Development',
-          ],
-        },
-      }
-    }
-    // Simulate boolean
-    if (schema?.answer) {
-      return {
-        object: { answer: 'true' },
-      }
-    }
-    // Default object
-    return { object: { result: 'Generated content' } }
-  }),
-  generateText: vi.fn().mockImplementation(async ({ prompt }) => {
-    return {
-      text: `Generated blog post for: ${prompt.slice(0, 50)}...`,
-    }
-  }),
-}))
+// Memory adapter for testing - simulates batch processing locally
+// Import from .ts file for proper vite resolution
+import { configureMemoryAdapter, clearBatches } from '../src/batch/memory.ts'
 
 // ============================================================================
 // Tests
@@ -84,7 +44,6 @@ vi.mock('../src/generate.js', () => ({
 
 describe('Implicit Batch Processing', () => {
   beforeEach(() => {
-    vi.clearAllMocks()
     resetContext()
     clearBatches()
     configureMemoryAdapter({})
@@ -145,7 +104,7 @@ describe('Implicit Batch Processing', () => {
     })
   })
 
-  describe('Three-Tier Execution (immediate → flex → batch)', () => {
+  describe('Three-Tier Execution (immediate -> flex -> batch)', () => {
     it('getExecutionTier returns immediate for < flexThreshold items', () => {
       configure({ batchMode: 'auto', flexThreshold: 5, batchThreshold: 500 })
 
@@ -267,8 +226,7 @@ describe('Implicit Batch Processing', () => {
 
       // Create batch map - this enters recording mode for each item
       const batchMap = createBatchMap(items, (item) => {
-        // When we call write` here, it should capture the operation
-        // Since we mocked generateText, we need to manually capture
+        // Capture operation for each item
         captureOperation(`Write about: ${item}`, 'text', undefined, undefined)
         recordedCount++
         return `result_${item}`
@@ -306,7 +264,7 @@ describe('Implicit Batch Processing', () => {
       })
     })
 
-    it('supports async iteration', async () => {
+    it('supports iteration over results', async () => {
       configure({ batchMode: 'immediate' })
 
       const items = ['X', 'Y']
@@ -334,20 +292,15 @@ describe('Implicit Batch Processing', () => {
   })
 
   describe('Full Workflow', () => {
-    it('list → map → batch flow works end-to-end', async () => {
+    it('list -> map -> batch flow works end-to-end', async () => {
       // Configure for immediate execution (for testing)
       configure({ batchMode: 'immediate', provider: 'openai', model: 'gpt-4o' })
 
-      // Step 1: Get titles (this executes immediately)
-      // Note: The mock returns { object: { items: [...] } }
-      // so we access .items from the result
-      const result = await list`5 blog post titles about startups`
-      const titles = (result as any).items || result
-      expect(titles).toHaveLength(5)
+      // Step 1: Simulate getting titles (in production this would be AI-generated)
+      const titles = ['Title 1', 'Title 2', 'Title 3', 'Title 4', 'Title 5']
 
       // Step 2: Map to blog posts
       // In the real implementation, this would capture operations
-      // For this test, we simulate the batch map behavior
       const batchMap = createBatchMap(titles, (title: string) => {
         // Capture the write operation
         captureOperation(`Write a blog post about: ${title}`, 'text')
@@ -384,12 +337,16 @@ describe('Implicit Batch Processing', () => {
       const items = ['Test']
       const batchMap = new BatchMapPromise<string>(
         items,
-        [[{
-          id: 'op_1',
-          prompt: 'Test prompt',
-          itemPlaceholder: 'Test',
-          type: 'text' as const,
-        }]],
+        [
+          [
+            {
+              id: 'op_1',
+              prompt: 'Test prompt',
+              itemPlaceholder: 'Test',
+              type: 'text' as const,
+            },
+          ],
+        ],
         { deferred: true }
       )