feat! : Support invoke with structured output in VercelAI provider

Author: jsonbailey

packages/ai-providers/server-ai-vercel/__tests__/VercelProvider.test.ts

@@ -1,10 +1,11 @@
-import { generateText } from 'ai';
+import { generateObject, generateText } from 'ai';
 
 import { VercelProvider } from '../src/VercelProvider';
 
 // Mock Vercel AI SDK
 jest.mock('ai', () => ({
   generateText: jest.fn(),
+  generateObject: jest.fn(),
 }));
 
 describe('VercelProvider', () => {
@@ -116,6 +117,76 @@ describe('VercelProvider', () => {
     });
   });
 
+  describe('convertToZodSchema', () => {
+    it('converts simple object structure to Zod schema', () => {
+      const responseStructure = {
+        name: 'string',
+        age: 0,
+        isActive: true,
+      };
+
+      const schema = VercelProvider.convertToZodSchema(responseStructure);
+
+      expect(schema).toBeDefined();
+      expect(typeof schema.parse).toBe('function');
+    });
+
+    it('converts nested object structure to Zod schema', () => {
+      const responseStructure = {
+        user: {
+          name: 'string',
+          age: 0,
+        },
+        settings: {
+          theme: 'string',
+          notifications: true,
+        },
+      };
+
+      const schema = VercelProvider.convertToZodSchema(responseStructure);
+
+      expect(schema).toBeDefined();
+      expect(typeof schema.parse).toBe('function');
+    });
+
+    it('converts array structure to Zod schema', () => {
+      const responseStructure = {
+        items: ['string'],
+        numbers: [0],
+        booleans: [true],
+      };
+
+      const schema = VercelProvider.convertToZodSchema(responseStructure);
+
+      expect(schema).toBeDefined();
+      expect(typeof schema.parse).toBe('function');
+    });
+
+    it('handles empty array structure', () => {
+      const responseStructure = {
+        items: [],
+      };
+
+      const schema = VercelProvider.convertToZodSchema(responseStructure);
+
+      expect(schema).toBeDefined();
+      expect(typeof schema.parse).toBe('function');
+    });
+
+    it('handles null and undefined values', () => {
+      const responseStructure = {
+        nullable: null,
+        undefined,
+        string: 'string',
+      };
+
+      const schema = VercelProvider.convertToZodSchema(responseStructure);
+
+      expect(schema).toBeDefined();
+      expect(typeof schema.parse).toBe('function');
+    });
+  });
+
   describe('invokeModel', () => {
     it('invokes Vercel AI generateText and returns response', async () => {
       const mockResponse = {
@@ -178,6 +249,91 @@ describe('VercelProvider', () => {
     });
   });
 
+  describe('invokeStructuredModel', () => {
+    it('invokes Vercel AI generateObject and returns structured response', async () => {
+      const mockResponse = {
+        object: {
+          name: 'John Doe',
+          age: 30,
+          isActive: true,
+        },
+        usage: {
+          promptTokens: 10,
+          completionTokens: 15,
+          totalTokens: 25,
+        },
+      };
+
+      (generateObject as jest.Mock).mockResolvedValue(mockResponse);
+
+      const messages = [{ role: 'user' as const, content: 'Generate user data' }];
+      const responseStructure = {
+        name: 'string',
+        age: 0,
+        isActive: true,
+      };
+
+      const result = await provider.invokeStructuredModel(messages, responseStructure);
+
+      expect(generateObject).toHaveBeenCalledWith({
+        model: mockModel,
+        messages: [{ role: 'user', content: 'Generate user data' }],
+        schema: expect.any(Object), // Zod schema
+      });
+
+      expect(result).toEqual({
+        data: {
+          name: 'John Doe',
+          age: 30,
+          isActive: true,
+        },
+        rawResponse: JSON.stringify({
+          name: 'John Doe',
+          age: 30,
+          isActive: true,
+        }),
+        metrics: {
+          success: true,
+          usage: {
+            total: 25,
+            input: 10,
+            output: 15,
+          },
+        },
+      });
+    });
+
+    it('handles structured response without usage data', async () => {
+      const mockResponse = {
+        object: {
+          result: 'success',
+        },
+      };
+
+      (generateObject as jest.Mock).mockResolvedValue(mockResponse);
+
+      const messages = [{ role: 'user' as const, content: 'Generate result' }];
+      const responseStructure = {
+        result: 'string',
+      };
+
+      const result = await provider.invokeStructuredModel(messages, responseStructure);
+
+      expect(result).toEqual({
+        data: {
+          result: 'success',
+        },
+        rawResponse: JSON.stringify({
+          result: 'success',
+        }),
+        metrics: {
+          success: true,
+          usage: undefined,
+        },
+      });
+    });
+  });
+
   describe('getModel', () => {
     it('returns the underlying Vercel AI model', () => {
       const model = provider.getModel();

packages/ai-providers/server-ai-vercel/package.json

@@ -48,7 +48,8 @@
     "jest": "^29.6.1",
     "prettier": "^3.0.0",
     "ts-jest": "^29.1.1",
-    "typescript": "5.1.6"
+    "typescript": "5.1.6",
+    "zod": "^3.22.0"
   },
   "peerDependencies": {
     "@ai-sdk/anthropic": "^2.0.0",
@@ -57,7 +58,8 @@
     "@ai-sdk/mistral": "^2.0.0",
     "@ai-sdk/openai": "^2.0.0",
     "@launchdarkly/server-sdk-ai": "^0.12.2",
-    "ai": "^4.0.0 || ^5.0.0"
+    "ai": "^4.0.0 || ^5.0.0",
+    "zod": "^3.22.0"
   },
   "peerDependenciesMeta": {
     "@ai-sdk/anthropic": {

packages/ai-providers/server-ai-vercel/src/VercelProvider.ts

@@ -1,4 +1,5 @@
-import { generateText, LanguageModel } from 'ai';
+import { generateObject, generateText, LanguageModel } from 'ai';
+import { z } from 'zod';
 
 import { AIProvider } from '@launchdarkly/server-sdk-ai';
 import type {
@@ -8,6 +9,7 @@ import type {
   LDLogger,
   LDMessage,
   LDTokenUsage,
+  StructuredResponse,
 } from '@launchdarkly/server-sdk-ai';
 
 /**
@@ -69,6 +71,34 @@ export class VercelProvider extends AIProvider {
     };
   }
 
+  /**
+   * Invoke the Vercel AI model with structured output support.
+   */
+  async invokeStructuredModel(
+    messages: LDMessage[],
+    responseStructure: Record<string, unknown>,
+  ): Promise<StructuredResponse> {
+    // Convert responseStructure to Zod schema
+    const schema = VercelProvider.convertToZodSchema(responseStructure);
+
+    // Call Vercel AI generateObject
+    const result = await generateObject({
+      model: this._model,
+      messages,
+      schema,
+      ...this._parameters,
+    });
+
+    // Extract metrics including token usage and success status
+    const metrics = VercelProvider.createAIMetrics(result);
+
+    return {
+      data: result.object,
+      rawResponse: JSON.stringify(result.object),
+      metrics,
+    };
+  }
+
   /**
    * Get the underlying Vercel AI model instance.
    */
@@ -121,6 +151,72 @@ export class VercelProvider extends AIProvider {
     };
   }
 
+  /**
+   * Convert a response structure object to a Zod schema.
+   * This method recursively converts a generic object structure to a Zod schema
+   * that can be used with Vercel AI SDK's generateObject function.
+   *
+   * @param responseStructure The structure object to convert
+   * @returns A Zod schema representing the structure
+   */
+  static convertToZodSchema(responseStructure: Record<string, unknown>): z.ZodSchema {
+    const shape: Record<string, z.ZodSchema> = {};
+
+    Object.entries(responseStructure).forEach(([key, value]) => {
+      // eslint-disable-next-line no-underscore-dangle
+      shape[key] = VercelProvider._convertValueToZodSchema(value);
+    });
+
+    return z.object(shape);
+  }
+
+  /**
+   * Convert a single value to a Zod schema.
+   * This is a helper method for convertToZodSchema that handles different value types.
+   *
+   * @param value The value to convert
+   * @returns A Zod schema for the value
+   */
+  private static _convertValueToZodSchema(value: unknown): z.ZodSchema {
+    if (value === null || value === undefined) {
+      return z.any();
+    }
+
+    if (typeof value === 'string') {
+      return z.string();
+    }
+
+    if (typeof value === 'number') {
+      return z.number();
+    }
+
+    if (typeof value === 'boolean') {
+      return z.boolean();
+    }
+
+    if (Array.isArray(value)) {
+      if (value.length === 0) {
+        return z.array(z.any());
+      }
+      // Use the first element to determine the array type
+      // eslint-disable-next-line no-underscore-dangle
+      const elementSchema = VercelProvider._convertValueToZodSchema(value[0]);
+      return z.array(elementSchema);
+    }
+
+    if (typeof value === 'object') {
+      const shape: Record<string, z.ZodSchema> = {};
+      Object.entries(value).forEach(([key, val]) => {
+        // eslint-disable-next-line no-underscore-dangle
+        shape[key] = VercelProvider._convertValueToZodSchema(val);
+      });
+      return z.object(shape);
+    }
+
+    // Fallback for any other type
+    return z.any();
+  }
+
   /**
    * Create a Vercel AI model from an AI configuration.
    * This method creates a Vercel AI model based on the provider configuration.