From dc04dbbff45f19e89033b27539ba729cc75e4030 Mon Sep 17 00:00:00 2001 From: Travis Fischer Date: Mon, 24 Mar 2025 17:18:49 +0800 Subject: [PATCH] feat: add support for ai-sdk-style execute param in createAIFunction --- packages/ai-sdk/src/ai-sdk.ts | 4 +- packages/core/src/create-ai-function.test.ts | 20 ++- packages/core/src/create-ai-function.ts | 129 +++++++++++++------ packages/core/src/schema.test.ts | 8 +- packages/core/src/schema.ts | 40 ++++-- packages/core/src/types.ts | 8 +- packages/genkit/src/genkit.ts | 7 +- packages/langchain/src/langchain.ts | 3 +- packages/llamaindex/src/llamaindex.ts | 9 +- packages/xsai/src/xsai.ts | 2 +- 10 files changed, 145 insertions(+), 85 deletions(-) diff --git a/packages/ai-sdk/src/ai-sdk.ts b/packages/ai-sdk/src/ai-sdk.ts index 01db994..f133a15 100644 --- a/packages/ai-sdk/src/ai-sdk.ts +++ b/packages/ai-sdk/src/ai-sdk.ts @@ -1,7 +1,7 @@ import { type AIFunctionLike, AIFunctionSet, - asSchema, + asAgenticSchema, isZodSchema } from '@agentic/core' import { jsonSchema, tool } from 'ai' @@ -20,7 +20,7 @@ export function createAISDKTools(...aiFunctionLikeTools: AIFunctionLike[]) { description: fn.spec.description, parameters: isZodSchema(fn.inputSchema) ? fn.inputSchema - : jsonSchema(asSchema(fn.inputSchema).jsonSchema), + : jsonSchema(asAgenticSchema(fn.inputSchema).jsonSchema), execute: fn.execute }) ]) diff --git a/packages/core/src/create-ai-function.test.ts b/packages/core/src/create-ai-function.test.ts index 0600a6e..cd08118 100644 --- a/packages/core/src/create-ai-function.test.ts +++ b/packages/core/src/create-ai-function.test.ts @@ -6,19 +6,17 @@ import { type Msg } from './message' // TODO: Add tests for passing JSON schema directly. -const fullNameAIFunction = createAIFunction( - { - name: 'fullName', - description: 'Returns the full name of a person.', - inputSchema: z.object({ - first: z.string(), - last: z.string() - }) - }, - async ({ first, last }) => { +const fullNameAIFunction = createAIFunction({ + name: 'fullName', + description: 'Returns the full name of a person.', + inputSchema: z.object({ + first: z.string(), + last: z.string() + }), + execute: ({ first, last }) => { return `${first} ${last}` } -) +}) describe('createAIFunction()', () => { test('exposes OpenAI function calling spec', () => { diff --git a/packages/core/src/create-ai-function.ts b/packages/core/src/create-ai-function.ts index 678259e..70acc28 100644 --- a/packages/core/src/create-ai-function.ts +++ b/packages/core/src/create-ai-function.ts @@ -1,7 +1,38 @@ import type * as types from './types' -import { asSchema } from './schema' +import { asAgenticSchema } from './schema' import { assert } from './utils' +export type CreateAIFunctionArgs< + InputSchema extends types.AIFunctionInputSchema +> = { + /** Name of the function. */ + name: string + + /** Description of the function. */ + description?: string + + /** + * Zod schema or AgenticSchema for the function parameters. + * + * You can use a JSON Schema for more dynamic tool sources such as MCP by + * using the `createJsonSchema` utility function. + */ + inputSchema: InputSchema + + /** + * Whether to enable strict structured output generation based on the given + * input schema. (this is a feature of the OpenAI API) + * + * Defaults to `true`. + */ + strict?: boolean +} + +export type AIFunctionImplementation< + InputSchema extends types.AIFunctionInputSchema, + Output +> = (params: types.inferInput) => types.MaybePromise + /** * Create a function meant to be used with OpenAI tool or function calling. * @@ -15,85 +46,99 @@ export function createAIFunction< InputSchema extends types.AIFunctionInputSchema, Output >( - spec: { - /** Name of the function. */ - name: string - /** Description of the function. */ - description?: string - /** Zod schema for the function parameters. */ - inputSchema: InputSchema - /** - * Whether or not to enable structured output generation based on the given - * zod schema. - */ - strict?: boolean + args: CreateAIFunctionArgs, + /** Underlying function implementation. */ + execute: AIFunctionImplementation +): types.AIFunction +export function createAIFunction< + InputSchema extends types.AIFunctionInputSchema, + Output +>( + args: CreateAIFunctionArgs & { + /** Underlying function implementation. */ + execute: AIFunctionImplementation + } +): types.AIFunction +export function createAIFunction< + InputSchema extends types.AIFunctionInputSchema, + Output +>( + { + name, + description = '', + inputSchema, + strict = true, + execute + }: CreateAIFunctionArgs & { + /** Underlying function implementation. */ + execute?: AIFunctionImplementation }, - /** Implementation of the function to call with the parsed arguments. */ - implementation: ( - params: types.inferInput - ) => types.MaybePromise + /** Underlying function implementation. */ + executeArg?: AIFunctionImplementation ): types.AIFunction { - assert(spec.name, 'createAIFunction missing required "spec.name"') + assert(name, 'createAIFunction missing required "name"') + assert(inputSchema, 'createAIFunction missing required "inputSchema"') assert( - spec.inputSchema, - 'createAIFunction missing required "spec.inputSchema"' + execute || executeArg, + 'createAIFunction missing required "execute" function implementation' ) - assert(implementation, 'createAIFunction missing required "implementation"') assert( - typeof implementation === 'function', - 'createAIFunction "implementation" must be a function' + !(execute && executeArg), + 'createAIFunction: cannot provide both "execute" and a second function argument. there should only be one function implementation.' + ) + execute ??= executeArg + assert( + execute, + 'createAIFunction missing required "execute" function implementation' + ) + assert( + typeof execute === 'function', + 'createAIFunction "execute" must be a function' ) - const strict = !!spec.strict - const inputSchema = asSchema(spec.inputSchema, { strict }) + const inputAgenticSchema = asAgenticSchema(inputSchema, { strict }) /** Parse the arguments string, optionally reading from a message. */ const parseInput = ( input: string | types.Msg ): types.inferInput => { if (typeof input === 'string') { - return inputSchema.parse(input) + return inputAgenticSchema.parse(input) } else { const args = input.function_call?.arguments assert( args, - `Missing required function_call.arguments for function ${spec.name}` + `Missing required function_call.arguments for function ${name}` ) - return inputSchema.parse(args) + return inputAgenticSchema.parse(args) } } - // Call the implementation function with the parsed arguments. + // Call the underlying function implementation with the parsed arguments. const aiFunction: types.AIFunction = ( input: string | types.Msg ) => { const parsedInput = parseInput(input) - return implementation(parsedInput) + return execute(parsedInput) } // Override the default function name with the intended name. Object.defineProperty(aiFunction, 'name', { - value: spec.name, + value: name, writable: false }) - aiFunction.inputSchema = spec.inputSchema + aiFunction.inputSchema = inputSchema aiFunction.parseInput = parseInput - + aiFunction.execute = execute aiFunction.spec = { - name: spec.name, - description: spec.description?.trim() ?? '', - parameters: inputSchema.jsonSchema, + name, + description, + parameters: inputAgenticSchema.jsonSchema, type: 'function', strict } - aiFunction.execute = ( - params: types.inferInput - ): types.MaybePromise => { - return implementation(params) - } - return aiFunction } diff --git a/packages/core/src/schema.test.ts b/packages/core/src/schema.test.ts index ae31c70..9b40099 100644 --- a/packages/core/src/schema.test.ts +++ b/packages/core/src/schema.test.ts @@ -1,18 +1,18 @@ import { expect, test } from 'vitest' import { z } from 'zod' -import { asSchema, createJsonSchema, isZodSchema } from './schema' +import { asAgenticSchema, createJsonSchema, isZodSchema } from './schema' test('isZodSchema', () => { expect(isZodSchema(z.object({}))).toBe(true) expect(isZodSchema({})).toBe(false) }) -test('asSchema', () => { - expect(asSchema(z.object({})).jsonSchema).toEqual({ +test('asAgenticSchema', () => { + expect(asAgenticSchema(z.object({})).jsonSchema).toEqual({ type: 'object', properties: {}, additionalProperties: false }) - expect(asSchema(createJsonSchema({})).jsonSchema).toEqual({}) + expect(asAgenticSchema(createJsonSchema({})).jsonSchema).toEqual({}) }) diff --git a/packages/core/src/schema.ts b/packages/core/src/schema.ts index 771626c..1bf8e88 100644 --- a/packages/core/src/schema.ts +++ b/packages/core/src/schema.ts @@ -10,7 +10,14 @@ import { zodToJsonSchema } from './zod-to-json-schema' */ export const schemaSymbol = Symbol('agentic.schema') -export type Schema = { +/** + * Structured schema used across Agentic, which wraps either a Zod schema or a + * JSON Schema. + * + * JSON Schema support is important to support more dynamic tool sources such as + * MCP. + */ +export type AgenticSchema = { /** * The JSON Schema. */ @@ -46,7 +53,7 @@ export type Schema = { _source?: any } -export function isSchema(value: unknown): value is Schema { +export function isAgenticSchema(value: unknown): value is AgenticSchema { return ( typeof value === 'object' && value !== null && @@ -59,24 +66,31 @@ export function isSchema(value: unknown): value is Schema { export function isZodSchema(value: unknown): value is z.ZodType { return ( + !!value && typeof value === 'object' && - value !== null && '_def' in value && '~standard' in value && - 'parse' in value && - 'safeParse' in value + (value['~standard'] as any)?.vendor === 'zod' ) } -export function asSchema( - schema: z.Schema | Schema, +export function asAgenticSchema( + schema: z.Schema | AgenticSchema, opts: { strict?: boolean } = {} -): Schema { - return isSchema(schema) ? schema : createSchemaFromZodSchema(schema, opts) +): AgenticSchema { + return isAgenticSchema(schema) + ? schema + : createAgenticSchemaFromZodSchema(schema, opts) +} + +export function asZodOrJsonSchema( + schema: z.Schema | AgenticSchema +): z.Schema | types.JSONSchema { + return isZodSchema(schema) ? schema : schema.jsonSchema } /** - * Create a Schema from a JSON Schema. + * Create an AgenticSchema from a JSON Schema. * * All `AIFunction` input schemas accept either a Zod schema or a custom JSON * Schema. Use this function to wrap JSON schemas for use with `AIFunction`. @@ -96,7 +110,7 @@ export function createJsonSchema( safeParse?: types.SafeParseFn source?: any } = {} -): Schema { +): AgenticSchema { safeParse ??= (value: unknown) => { try { const result = parse(value) @@ -116,10 +130,10 @@ export function createJsonSchema( } } -export function createSchemaFromZodSchema( +export function createAgenticSchemaFromZodSchema( zodSchema: z.Schema, opts: { strict?: boolean } = {} -): Schema { +): AgenticSchema { return createJsonSchema(zodToJsonSchema(zodSchema, opts), { parse: (value) => { return parseStructuredOutput(value, zodSchema) diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index d1c4b57..6bd1b9b 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -4,10 +4,10 @@ import type { z } from 'zod' import type { AIFunctionSet } from './ai-function-set' import type { AIFunctionsProvider } from './fns' import type { Msg } from './message' -import type { Schema } from './schema' +import type { AgenticSchema } from './schema' export type { Msg } from './message' -export type { Schema } from './schema' +export type { AgenticSchema } from './schema' export type { KyInstance } from 'ky' export type { ThrottledFunction } from 'p-throttle' export type { SetOptional, SetRequired, Simplify } from 'type-fest' @@ -57,11 +57,11 @@ export interface AIToolSpec { * A Zod object schema or a custom schema created from a JSON schema via * `createSchema()`. */ -export type AIFunctionInputSchema = z.ZodObject | Schema +export type AIFunctionInputSchema = z.ZodObject | AgenticSchema // eslint-disable-next-line @typescript-eslint/naming-convention export type inferInput = - InputSchema extends Schema + InputSchema extends AgenticSchema ? InputSchema['_type'] : InputSchema extends z.ZodTypeAny ? z.infer diff --git a/packages/genkit/src/genkit.ts b/packages/genkit/src/genkit.ts index 04b7f89..e9d3fb9 100644 --- a/packages/genkit/src/genkit.ts +++ b/packages/genkit/src/genkit.ts @@ -2,7 +2,7 @@ import type { Genkit } from 'genkit' import { type AIFunctionLike, AIFunctionSet, - asSchema, + asZodOrJsonSchema, isZodSchema } from '@agentic/core' import { z } from 'zod' @@ -26,10 +26,7 @@ export function createGenkitTools( { name: fn.spec.name, description: fn.spec.description, - // TODO: This schema handling should be able to be cleaned up. - [inputSchemaKey]: isZodSchema(fn.inputSchema) - ? fn.inputSchema - : asSchema(fn.inputSchema).jsonSchema, + [inputSchemaKey]: asZodOrJsonSchema(fn.inputSchema), outputSchema: z.any() }, fn.execute diff --git a/packages/langchain/src/langchain.ts b/packages/langchain/src/langchain.ts index e908bf3..3add090 100644 --- a/packages/langchain/src/langchain.ts +++ b/packages/langchain/src/langchain.ts @@ -1,6 +1,7 @@ import { type AIFunctionLike, AIFunctionSet, + asZodOrJsonSchema, stringifyForModel } from '@agentic/core' import { DynamicStructuredTool } from '@langchain/core/tools' @@ -17,7 +18,7 @@ export function createLangChainTools(...aiFunctionLikeTools: AIFunctionLike[]) { new DynamicStructuredTool({ name: fn.spec.name, description: fn.spec.description, - schema: fn.inputSchema, + schema: asZodOrJsonSchema(fn.inputSchema), func: async (input) => { const result = await Promise.resolve(fn.execute(input)) // LangChain tools require the output to be a string diff --git a/packages/llamaindex/src/llamaindex.ts b/packages/llamaindex/src/llamaindex.ts index bf3a706..7c85068 100644 --- a/packages/llamaindex/src/llamaindex.ts +++ b/packages/llamaindex/src/llamaindex.ts @@ -1,4 +1,8 @@ -import { type AIFunctionLike, AIFunctionSet } from '@agentic/core' +import { + type AIFunctionLike, + AIFunctionSet, + asZodOrJsonSchema +} from '@agentic/core' import { FunctionTool } from 'llamaindex' /** @@ -14,7 +18,8 @@ export function createLlamaIndexTools( FunctionTool.from(fn.execute, { name: fn.spec.name, description: fn.spec.description, - parameters: fn.spec.parameters as any + // TODO: Investigate types here + parameters: asZodOrJsonSchema(fn.inputSchema) as any }) ) } diff --git a/packages/xsai/src/xsai.ts b/packages/xsai/src/xsai.ts index 13fbbb9..be78b8d 100644 --- a/packages/xsai/src/xsai.ts +++ b/packages/xsai/src/xsai.ts @@ -3,7 +3,7 @@ import { tool, type ToolResult } from '@xsai/tool' /** * Converts a set of Agentic stdlib AI functions to an object compatible with - * [the xsAI SDK's](https://github.com/moeru-ai/xsai) `tools` parameter. + * the [xsAI SDK's](https://github.com/moeru-ai/xsai) `tools` parameter. */ export function createXSAITools( ...aiFunctionLikeTools: AIFunctionLike[]