Merge pull request #4 from petrbrzek/add-support-for-array-type

Add Support for Array Types
hmarr · Aug 7, 2023 · fd29510 · fd29510
2 parents 426712b + cc17cc8
commit fd29510
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/src/functions.ts b/src/functions.ts
@@ -31,6 +31,10 @@ type Prop = {
     }
     | { type: "boolean" }
     | { type: "null" }
+    | {
+        type: "array";
+        items?: Prop;
+      }
   );
 
 // When OpenAI use functions in the prompt, they format them as TypeScript definitions rather than OpenAPI JSON schemas.
@@ -94,5 +98,10 @@ function formatType(param: Prop, indent: number): string {
       return "null";
     case "object":
       return ["{", formatObjectProperties(param, indent + 2), "}"].join("\n");
+    case "array":
+      if (param.items) {
+        return `${formatType(param.items, indent)}[]`;
+      }
+      return "any[]";
   }
 }
diff --git a/tests/token-counts.test.ts b/tests/token-counts.test.ts
@@ -215,6 +215,51 @@ const TEST_CASES: Example[] = [
     ],
     tokens: 40,
   },
+  {
+    messages: [{ role: 'user', content: 'hello' }],
+    functions: [
+      {
+        name: 'get_recipe',
+        parameters: {
+          type: 'object',
+          required: ['ingredients', 'instructions', 'time_to_cook'],
+          properties: {
+            ingredients: {
+              type: 'array',
+              items: {
+                type: 'object',
+                required: ['name', 'unit', 'amount'],
+                properties: {
+                  name: {
+                    type: 'string',
+                  },
+                  unit: {
+                    enum: ['grams', 'ml', 'cups', 'pieces', 'teaspoons'],
+                    type: 'string',
+                  },
+                  amount: {
+                    type: 'number',
+                  },
+                },
+              },
+            },
+            instructions: {
+              type: 'array',
+              items: {
+                type: 'string',
+              },
+              description: 'Steps to prepare the recipe (no numbering)',
+            },
+            time_to_cook: {
+              type: 'number',
+              description: 'Total time to prepare the recipe in minutes',
+            },
+          },
+        },
+      },
+    ],
+    tokens: 106,
+  },
 ];
 
 const validateAll = false;