feat: support anyOf in function parameters (#13)

hmarr · Sep 28, 2023 · 1ab24fe · 1ab24fe
1 parent 5873db7
commit 1ab24fe
Show file tree

Hide file tree

Showing 2 changed files with 169 additions and 0 deletions.
diff --git a/src/functions.ts b/src/functions.ts
@@ -19,9 +19,14 @@ interface ObjectProp {
   required?: string[];
 }
 
+interface AnyOfProp {
+  anyOf: Prop[];
+}
+
 type Prop = {
   description?: string;
 } & (
+  | AnyOfProp
   | ObjectProp
   | {
       type: "string";
@@ -41,6 +46,13 @@ type Prop = {
     }
 );
 
+function isAnyOfProp(prop: Prop): prop is AnyOfProp {
+  return (
+    (prop as AnyOfProp).anyOf !== undefined &&
+    Array.isArray((prop as AnyOfProp).anyOf)
+  );
+}
+
 // When OpenAI use functions in the prompt, they format them as TypeScript definitions rather than OpenAPI JSON schemas.
 // This function converts the JSON schemas into TypeScript definitions.
 export function formatFunctionDefinitions(functions: FunctionDef[]) {
@@ -80,6 +92,9 @@ function formatObjectProperties(obj: ObjectProp, indent: number): string {
 
 // Format a single property type
 function formatType(param: Prop, indent: number): string {
+  if (isAnyOfProp(param)) {
+    return param.anyOf.map((v) => formatType(v, indent)).join(" | ");
+  }
   switch (param.type) {
     case "string":
       if (param.enum) {

diff --git a/tests/token-counts.test.ts b/tests/token-counts.test.ts
@@ -264,6 +264,160 @@ const TEST_CASES: Example[] = [
     ],
     tokens: 25,
   },
+  {
+    messages: [
+      {
+        role: "system",
+        content: "You are an AI assistant trained to foo or bar",
+      },
+      { role: "user", content: "My name is suzie" },
+      {
+        role: "function",
+        name: "foo",
+        content: '{"res":{"kind":"user","name":"suzie"}}',
+      },
+      {
+        role: "user",
+        content: 'I want to post a message with the text "hello world"',
+      },
+      {
+        role: "function",
+        name: "foo",
+        content: '{"res":{"kind":"post","text":"hello world"}}',
+      },
+    ],
+    functions: [
+      {
+        name: "foo",
+        description: "Return the foo or the bar",
+        parameters: {
+          type: "object",
+          properties: {
+            res: {
+              anyOf: [
+                {
+                  type: "object",
+                  properties: {
+                    kind: { type: "string", const: "post" },
+                    text: { type: "string" },
+                  },
+                  required: ["kind", "text"],
+                  additionalProperties: false,
+                },
+                {
+                  type: "object",
+                  properties: {
+                    kind: { type: "string", const: "user" },
+                    name: {
+                      type: "string",
+                      enum: ["jane", "suzie", "adam"],
+                    },
+                  },
+                  required: ["kind", "name"],
+                  additionalProperties: false,
+                },
+              ],
+              description: "The foo or the bar",
+            },
+          },
+          required: ["res"],
+          additionalProperties: false,
+        },
+      },
+    ],
+    function_call: { name: "foo" },
+    tokens: 158,
+  },
+  {
+    messages: [
+      { role: "system", content: "Hello" },
+      { role: "user", content: "Hi there" },
+    ],
+    functions: [
+      {
+        name: "do_stuff",
+        parameters: {
+          type: "object",
+          properties: {
+            foo: {
+              anyOf: [
+                {
+                  type: "object",
+                  properties: {
+                    kind: { type: "string", const: "a" },
+                    baz: { type: "boolean" },
+                  },
+                },
+              ],
+            },
+          },
+        },
+      },
+    ],
+    tokens: 52,
+  },
+  {
+    messages: [
+      { role: "system", content: "Hello" },
+      { role: "user", content: "Hi there" },
+    ],
+    functions: [
+      {
+        name: "do_stuff",
+        parameters: {
+          type: "object",
+          properties: {
+            foo: {
+              anyOf: [
+                {
+                  type: "object",
+                  properties: {
+                    kind: { type: "string", const: "a" },
+                    baz: { type: "boolean" },
+                  },
+                },
+                {
+                  type: "object",
+                  properties: {
+                    kind: { type: "string", const: "b" },
+                    qux: { type: "number" },
+                  },
+                },
+                {
+                  type: "object",
+                  properties: {
+                    kind: { type: "string", const: "c" },
+                    idk: { type: "string" },
+                  },
+                },
+              ],
+            },
+          },
+        },
+      },
+    ],
+    tokens: 80,
+  },
+  {
+    messages: [
+      { role: "system", content: "Hello" },
+      { role: "user", content: "Hi there" },
+    ],
+    functions: [
+      {
+        name: "do_stuff",
+        parameters: {
+          type: "object",
+          properties: {
+            foo: {
+              anyOf: [{ type: "string", const: "a" }, { type: "number" }],
+            },
+          },
+        },
+      },
+    ],
+    tokens: 44,
+  },
   {
     messages: [
       { role: "system", content: "Hello" },