Merge pull request #597 from mkbehr/structured-output-enum

Add enum support to structured outputs.

Author: vivtsai

functions/src/api/gemini.api.test.ts

@@ -98,6 +98,14 @@ describe('Gemini API', () => {
               description: 'An integer-valued property',
             },
           },
+          {
+            name: 'enumProperty',
+            schema: {
+              type: StructuredOutputDataType.ENUM,
+              description: 'An enum-valued property',
+              enumItems: ['FOO', 'BAR', 'BAZ'],
+            },
+          },
         ],
       },
     };
@@ -126,9 +134,14 @@ describe('Gemini API', () => {
               type: 'INTEGER',
               description: 'An integer-valued property',
             },
+            enumProperty: {
+              type: 'STRING',
+              description: 'An enum-valued property',
+              enum: ['FOO', 'BAR', 'BAZ'],
+            },
           },
-          propertyOrdering: ['stringProperty', 'integerProperty'],
-          required: ['stringProperty', 'integerProperty'],
+          propertyOrdering: ['stringProperty', 'integerProperty', 'enumProperty'],
+          required: ['stringProperty', 'integerProperty', 'enumProperty'],
         },
       },
     };

functions/src/api/gemini.api.ts

@@ -48,6 +48,7 @@ function makeStructuredOutputSchema(schema: StructuredOutputSchema): object {
     [StructuredOutputDataType.BOOLEAN]: 'BOOLEAN',
     [StructuredOutputDataType.ARRAY]: 'ARRAY',
     [StructuredOutputDataType.OBJECT]: 'OBJECT',
+    [StructuredOutputDataType.ENUM]: 'STRING',
   };
   const type = typeMap[schema.type];
   if (!type) {
@@ -79,6 +80,7 @@ function makeStructuredOutputSchema(schema: StructuredOutputSchema): object {
     properties: properties,
     propertyOrdering: orderedPropertyNames,
     required: orderedPropertyNames,
+    enum: schema.enumItems,
     items: itemsSchema,
   };
 }

functions/src/api/openai.api.test.ts

@@ -115,6 +115,14 @@ describe('OpenAI-compatible API', () => {
                 description: 'An integer-valued property',
               },
             },
+            {
+            name: 'enumProperty',
+            schema: {
+              type: StructuredOutputDataType.ENUM,
+              description: 'An enum-valued property',
+              enumItems: ['FOO', 'BAR', 'BAZ'],
+            },
+          },
           ],
         },
       });
@@ -147,9 +155,14 @@ describe('OpenAI-compatible API', () => {
           type: 'INTEGER',
           description: 'An integer-valued property',
         },
+        enumProperty: {
+              type: 'STRING',
+              description: 'An enum-valued property',
+              enum: ['FOO', 'BAR', 'BAZ'],
+        },
       },
       additionalProperties: false,
-      required: ['stringProperty', 'integerProperty'],
+      required: ['stringProperty', 'integerProperty', 'enumProperty'],
     };
     expect(parsedResponse.response_format).toEqual({
       type: 'json_schema',

functions/src/api/openai.api.ts

@@ -22,6 +22,7 @@ function makeStructuredOutputSchema(schema: StructuredOutputSchema): object {
     [StructuredOutputDataType.BOOLEAN]: 'BOOLEAN',
     [StructuredOutputDataType.ARRAY]: 'ARRAY',
     [StructuredOutputDataType.OBJECT]: 'OBJECT',
+    [StructuredOutputDataType.ENUM]: 'STRING',
   };
   const type = typeMap[schema.type];
   if (!type) {
@@ -55,6 +56,7 @@ function makeStructuredOutputSchema(schema: StructuredOutputSchema): object {
     properties: properties,
     additionalProperties: additionalProperties,
     required: orderedPropertyNames,
+    enum: schema.enumItems,
     items: itemsSchema,
   };
 }

utils/src/structured_output.test.ts

@@ -23,6 +23,14 @@ describe('Structured outputs', () => {
             },
           },
         },
+        {
+          name: 'enumProperty',
+          schema: {
+            type: StructuredOutputDataType.ENUM,
+            description: 'An enum-valued property',
+            enumItems: ['FOO', 'BAR', 'BAZ'],
+          },
+        },
       ],
     };
 
@@ -43,8 +51,13 @@ describe('Structured outputs', () => {
             type: 'integer',
           },
         },
+        enumProperty: {
+          description: 'An enum-valued property',
+          type: 'enum',
+          enum: ['FOO', 'BAR', 'BAZ'],
+        },
       },
-      required: ['stringProperty', 'intArrayProperty'],
+      required: ['stringProperty', 'intArrayProperty', 'enumProperty'],
     };
     expect(parsedResult).toEqual(expectedResult);
   });

utils/src/structured_output.ts

@@ -16,13 +16,15 @@ export enum StructuredOutputDataType {
   BOOLEAN = 'BOOLEAN',
   ARRAY = 'ARRAY',
   OBJECT = 'OBJECT',
+  ENUM = 'ENUM',
 }
 
 export interface StructuredOutputSchema {
   type: StructuredOutputDataType;
   description?: string;
   properties?: {name: string; schema: StructuredOutputSchema}[];
   arrayItems?: StructuredOutputSchema;
+  enumItems?: string[];
 }
 
 // TODO: Add StageKind or new enum to differentiate different configs
@@ -123,6 +125,7 @@ function schemaToObject(schema: StructuredOutputSchema): object {
     type: schema.type.toLowerCase(),
     properties: properties ?? undefined,
     items: arrayItems ?? undefined,
+    enum: schema.enumItems ?? undefined,
     required: required,
   };
 }