Change function parameters to be JSON Schema (#536)

Previously the function definitions were typed as either Zod or JSON Schema-ish (with some differences around `required`). Rather than have a layer of conversion logic (`getParametersSchema`) this makes the input types explicitly JSON Schema. This facilitates integration with other types, e.g. OpenAPI specs. (Callers using Zod must now do the conversion themselves, e.g. with `zod-to-json-schema`.)
fixie-ai · Jan 22, 2024 · 4ce9b17 · 4ce9b17 · vercel · Jan 22, 2024
1 parent af8be60
commit 4ce9b17
Show file tree

Hide file tree

Showing 17 changed files with 249 additions and 268 deletions.
diff --git a/packages/ai-jsx/package.json b/packages/ai-jsx/package.json
@@ -4,7 +4,7 @@
   "repository": "fixie-ai/ai-jsx",
   "bugs": "https://github.com/fixie-ai/ai-jsx/issues",
   "homepage": "https://ai-jsx.com",
-  "version": "0.28.2",
+  "version": "0.29.0",
   "volta": {
     "extends": "../../package.json"
   },
@@ -381,6 +381,7 @@
     "@opentelemetry/api": "^1.4.1",
     "@opentelemetry/api-logs": "^0.41.1",
     "@ozymandiasthegreat/vad": "^2.0.7",
+    "@types/json-schema": "^7.0.15",
     "axios": "^1.4.0",
     "cli-highlight": "^2.1.11",
     "cli-spinners": "^2.9.0",

diff --git a/packages/ai-jsx/src/batteries/docs.tsx b/packages/ai-jsx/src/batteries/docs.tsx
@@ -669,11 +669,14 @@ export class FixieCorpus<ChunkMetadata extends Jsonifiable = Jsonifiable> implem
     return {
       description,
       parameters: {
-        query: {
-          description: 'The search query. It will be embedded and used in a vector search against the corpus.',
-          type: 'string',
-          required: true,
+        type: 'object',
+        properties: {
+          query: {
+            description: 'The search query. It will be embedded and used in a vector search against the corpus.',
+            type: 'string',
+          },
         },
+        required: ['query'],
       },
       func: async function FixieCorpusQuery({ query }: { query: string }, { getContext }) {
         const corpus = new FixieCorpus(corpusId, getContext(FixieAPIContext));

diff --git a/packages/ai-jsx/src/batteries/sidekick/platform/large-response-handler.tsx b/packages/ai-jsx/src/batteries/sidekick/platform/large-response-handler.tsx
@@ -189,11 +189,14 @@ export function redactedFunctionTools(messages: ConversationMessage[]): UseTools
     loadBySimilarity: {
       description: 'Query the response of the "latest redacted function call" by using semantic similarity search.',
       parameters: {
-        query: {
-          type: 'string',
-          description: 'A query string.',
-          required: true,
+        type: 'object',
+        properties: {
+          query: {
+            type: 'string',
+            description: 'A query string.',
+          },
         },
+        required: ['query'],
       },
       func: ({ query }) => (
         <RerankerFormatted

diff --git a/packages/ai-jsx/src/batteries/use-tools.tsx b/packages/ai-jsx/src/batteries/use-tools.tsx
@@ -4,25 +4,15 @@
  * @packageDocumentation
  */
 
-import { ChatCompletion, FunctionParameters, FunctionResponse } from '../core/completion.js';
+import { ChatCompletion, FunctionDefinition, FunctionResponse } from '../core/completion.js';
 import { Component, ComponentContext, Node, RenderContext } from '../index.js';
 import { ConversationMessage, Converse, renderToConversation } from '../core/conversation.js';
 import _ from 'lodash';
 
 /**
  * Represents a tool that can be provided for the Large Language Model.
  */
-export interface Tool {
-  /**
-   * A description of what the tool does.
-   */
-  description: string;
-
-  /**
-   * A map of parameter names to their description and type.
-   */
-  parameters: FunctionParameters;
-
+export interface Tool extends FunctionDefinition {
   /**
    * A function that invokes the tool.
    *
@@ -32,7 +22,6 @@ export interface Tool {
    * can return a `string` or any AI.JSX {@link Node}, synchronously or
    * asynchronously.
    */
-  // Can we use Zod to do better than any?
   func: Component<any>;
 }
 
@@ -54,13 +43,6 @@ export interface UseToolsProps {
    * Whether the result should include intermediate steps, for example, the execution of the function and its response.
    */
   showSteps?: boolean;
-
-  /**
-   * User data the AI can use to determine what parameters to invoke the tool with.
-   *
-   * For instance, if the user's query can be "what's the weather like at my current location", you might pass `userData` as { "location": "Seattle" }.
-   */
-  userData?: string;
 }
 
 /**
@@ -102,7 +84,7 @@ export async function ExecuteFunction<T>(
  *  // Activate a scene in the user's lighting settings, like "Bedtime" or "Midday".
  *  async function activeScene({sceneName}: {sceneName: string}) { ... Code to activate a scene ... }
  *
- *  import z from 'zod';
+
  *  const tools: Record<string, Tool> = {
  *    turnLightsOn: {
  *      description: "Turn the lights on in the user's home",
@@ -117,11 +99,14 @@ export async function ExecuteFunction<T>(
  *    activeScene: {
  *      description: `Activate a scene in the user's lighting settings, like "Bedtime" or "Midday".`,
  *      parameters: {
- *        sceneName: {
- *          description: "The scene to activate the lighting in.",
- *          type: "string",
- *          required: true,
+ *        type: "object",
+ *        properties: {
+ *          sceneName: {
+ *            description: "The scene to activate the lighting in.",
+ *            type: "string",
+ *           },
  *        },
+ *        required: ["sceneName"],
  *      },
  *      func: activeScene,
  *    },

diff --git a/packages/ai-jsx/src/core/completion.tsx b/packages/ai-jsx/src/core/completion.tsx
@@ -9,9 +9,7 @@ import { AIJSXError, ErrorCode } from '../core/errors.js';
 import { OpenAIChatModel, OpenAICompletionModel } from '../lib/openai.js';
 import { getEnvVar } from '../lib/util.js';
 import { AnthropicChatModel } from '../lib/anthropic.js';
-import z from 'zod';
-import { zodToJsonSchema } from 'zod-to-json-schema';
-import _ from 'lodash';
+import { type JSONSchema7 } from 'json-schema';
 export {
   UserMessage,
   SystemMessage,
@@ -58,68 +56,9 @@ export type ModelComponent<T extends ModelPropsWithChildren> = Component<T>;
  */
 export interface FunctionDefinition {
   description?: string;
-  parameters: FunctionParameters;
+  parameters: JSONSchema7 & { type?: 'object' };
 }
 
-/**
- * This function creates a [JSON Schema](https://json-schema.org/) object to describe
- * parameters for a {@link FunctionDefinition}.
- *
- * See {@link FunctionParameters} for more information on what parameters are supported.
- */
-export function getParametersSchema(parameters: FunctionParameters) {
-  if (parameters instanceof z.ZodObject) {
-    return zodToJsonSchema(parameters);
-  }
-  return {
-    type: 'object',
-    required: Object.keys(parameters).filter((name) => parameters[name].required),
-    properties: Object.keys(parameters).reduce(
-      (map: Record<string, any>, paramName) => ({
-        ...map,
-        [paramName]: _.omit(parameters[paramName], 'required'),
-      }),
-      {}
-    ),
-  };
-}
-
-/**
- * Represents parameters to a {@link FunctionDefinition}.
- *
- * This is a simplified version of the `parameters` field in {@link ChatCompletionFunctions}: https://platform.openai.com/docs/api-reference/chat/create#chat/create-parameters.
- *
- *
- * If you want to pass a field to {@link FunctionParameters} that isn't supported on this type, you can use a {@link z.ZodObject} schema instead.
- */
-export interface PlainFunctionParameter {
-  description?: string;
-  type?: string;
-  /**
-   * The possible values this param can take.
-   */
-  enum?: string[];
-  required: boolean;
-}
-
-/**
- * Represents parameters to a {@link FunctionDefinition}.
- *
- * This type allows two ways for specifying parameters:
- * - For simple use cases, a record of parameter names to {@link PlainFunctionParameter} objects.
- * - For more complex use cases, a {@link z.ZodObject} schema object (`zod` is a standard runtime type definition & checking library).
- *
- * @note If using a Zod schema, the top-level schema must be an object as per OpenAI specifications:
- * https://platform.openai.com/docs/api-reference/chat/create#chat/create-parameters
- *
- * For example, to describe a list of strings, the following is not accepted:
- * `const schema: z.Schema = z.array(z.string())`
- *
- * Instead, you can wrap it in an object like so:
- * `const schema: z.ZodObject = z.object({ arr: z.array(z.string()) })`
- */
-export type FunctionParameters = Record<string, PlainFunctionParameter> | z.ZodObject<any>;
-
 /**
  * If env var `OPENAI_API_KEY` is defined, use Open AI as the completion model provider.
  *

diff --git a/packages/ai-jsx/src/lib/openai.tsx b/packages/ai-jsx/src/lib/openai.tsx
@@ -10,7 +10,6 @@ import {
   FunctionDefinition,
   ModelProps,
   ModelPropsWithChildren,
-  getParametersSchema,
 } from '../core/completion.js';
 import { AssistantMessage, ConversationMessage, FunctionCall, renderToConversation } from '../core/conversation.js';
 import { AIJSXError, ErrorCode } from '../core/errors.js';
@@ -508,7 +507,7 @@ export async function* OpenAIChatModel(
           function: {
             name: functionName,
             description: functionDefinition.description,
-            parameters: getParametersSchema(functionDefinition.parameters),
+            parameters: functionDefinition.parameters as Record<string, unknown>,
           },
           type: 'function',
         })

diff --git a/packages/docs/docs/ai-newcomers.md b/packages/docs/docs/ai-newcomers.md
@@ -110,19 +110,21 @@ In AI.JSX, this looks like:
 /**
  * Activate a scene in the user's lighting settings, like "Bedtime" or "Midday".
  */
-async function activateScene({sceneName}: {sceneName: string}) {}
+async function activateScene({ sceneName }: { sceneName: string }) {}
 
 // Describe to the LLM what's available
-import z from 'zod';
 const tools: Record<string, Tool> = {
   activateScene: {
     description: `Activate a scene in the user's lighting settings, like "Bedtime" or "Midday".`,
-    parameters: parameters: {
-      sceneName: {
-        description: "The scene to activate the lighting in.",
-        type: "string",
-        required: true,
+    parameters: {
+      type: 'object',
+      properties: {
+        sceneName: {
+          description: 'The scene to activate the lighting in.',
+          type: 'string',
+        },
       },
+      required: ['sceneName'],
     },
     func: activateScene,
   },

diff --git a/packages/docs/docs/changelog.md b/packages/docs/docs/changelog.md
@@ -1,6 +1,11 @@
 # Changelog
 
-## 0.28.2
+## 0.30.0
+
+- Changed the `FunctionDefinition` and `Tool` types to explicit JSON Schema. Zod types must now be explicitly
+  converted to JSON Schema, and the `required` semantics now match JSON Schema.
+
+## [0.28.2](https://github.com/fixie-ai/ai-jsx/tree/524cf2d69eb63426a098723997e39f0dda5a37c2)
 
 - Fix bug where partially streamed unicode characters (e.g. Chinese) would cause an error in OpenAI function calls.
 

diff --git a/packages/docs/docs/tutorials/aijsxTutorials/part7-tools.md b/packages/docs/docs/tutorials/aijsxTutorials/part7-tools.md
@@ -21,22 +21,28 @@ const tools = {
   checkStockPrice: {
     description: 'Check the price of a stock.',
     parameters: {
-      symbol: {
-        description: 'The symbol of the stock to get price for.',
-        type: 'string',
-        required: true,
+      type: 'object',
+      properties: {
+        symbol: {
+          description: 'The symbol of the stock to get price for.',
+          type: 'string',
+        },
       },
+      required: ['symbol'],
     },
     func: checkStockPrice,
   },
   getHistoricalPrices: {
     description: 'Return historical prices for a stock.',
     parameters: {
-      symbol: {
-        description: 'The symbol of the stock to get price for.',
-        type: 'string',
-        required: true,
+      type: 'object',
+      properties: {
+        symbol: {
+          description: 'The symbol of the stock to get price for.',
+          type: 'string',
+        },
       },
+      required: ['symbol'],
     },
     func: getHistoricalPrices,
   },