feat: add support for ai-sdk-style execute param in createAIFunction

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
       })
     ])

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(
+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 }) => {
-  async ({ first, last }) => {
     return `${first} ${last}`
   }
-)
+})
 
 describe('createAIFunction()', () => {
   test('exposes OpenAI function calling spec', () => {

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<InputSchema>) => types.MaybePromise<Output>
+
 /**
  * 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: {
+  args: CreateAIFunctionArgs<InputSchema>,
-    /** Name of the function. */
+  /** Underlying function implementation. */
-    name: string
+  execute: AIFunctionImplementation<InputSchema, Output>
-    /** Description of the function. */
+): types.AIFunction<InputSchema, Output>
-    description?: string
+export function createAIFunction<
-    /** Zod schema for the function parameters. */
+  InputSchema extends types.AIFunctionInputSchema,
-    inputSchema: InputSchema
+  Output
-    /**
+>(
-     * Whether or not to enable structured output generation based on the given
+  args: CreateAIFunctionArgs<InputSchema> & {
-     * zod schema.
+    /** Underlying function implementation. */
-     */
+    execute: AIFunctionImplementation<InputSchema, Output>
-    strict?: boolean
+  }
+): types.AIFunction<InputSchema, Output>
+export function createAIFunction<
+  InputSchema extends types.AIFunctionInputSchema,
+  Output
+>(
+  {
+    name,
+    description = '',
+    inputSchema,
+    strict = true,
+    execute
+  }: CreateAIFunctionArgs<InputSchema> & {
+    /** Underlying function implementation. */
+    execute?: AIFunctionImplementation<InputSchema, Output>
   },
-  /** Implementation of the function to call with the parsed arguments. */
+  /** Underlying function implementation. */
-  implementation: (
+  executeArg?: AIFunctionImplementation<InputSchema, Output>
-    params: types.inferInput<InputSchema>
-  ) => types.MaybePromise<Output>
 ): types.AIFunction<InputSchema, Output> {
-  assert(spec.name, 'createAIFunction missing required "spec.name"')
+  assert(name, 'createAIFunction missing required "name"')
+  assert(inputSchema, 'createAIFunction missing required "inputSchema"')
   assert(
-    spec.inputSchema,
+    execute || executeArg,
-    'createAIFunction missing required "spec.inputSchema"'
+    'createAIFunction missing required "execute" function implementation'
   )
-  assert(implementation, 'createAIFunction missing required "implementation"')
   assert(
-    typeof implementation === 'function',
+    !(execute && executeArg),
-    'createAIFunction "implementation" must be a function'
+    '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 inputAgenticSchema = asAgenticSchema(inputSchema, { strict })
-  const inputSchema = asSchema(spec.inputSchema, { strict })
 
   /** Parse the arguments string, optionally reading from a message. */
   const parseInput = (
     input: string | types.Msg
   ): types.inferInput<InputSchema> => {
     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<InputSchema, Output> = (
     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,
+    name,
-    description: spec.description?.trim() ?? '',
+    description,
-    parameters: inputSchema.jsonSchema,
+    parameters: inputAgenticSchema.jsonSchema,
     type: 'function',
     strict
   }
 
-  aiFunction.execute = (
-    params: types.inferInput<InputSchema>
-  ): types.MaybePromise<Output> => {
-    return implementation(params)
-  }
-
   return aiFunction
 }

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', () => {
+test('asAgenticSchema', () => {
-  expect(asSchema(z.object({})).jsonSchema).toEqual({
+  expect(asAgenticSchema(z.object({})).jsonSchema).toEqual({
     type: 'object',
     properties: {},
     additionalProperties: false
   })
-  expect(asSchema(createJsonSchema({})).jsonSchema).toEqual({})
+  expect(asAgenticSchema(createJsonSchema({})).jsonSchema).toEqual({})
 })

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<TData = unknown> = {
+/**
+ * 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<TData = unknown> = {
   /**
    * The JSON Schema.
    */
@@ -46,7 +53,7 @@ export type Schema<TData = unknown> = {
   _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 &&
+    (value['~standard'] as any)?.vendor === 'zod'
-    'safeParse' in value
   )
 }
 
-export function asSchema<TData>(
+export function asAgenticSchema<TData>(
-  schema: z.Schema<TData> | Schema<TData>,
+  schema: z.Schema<TData> | AgenticSchema<TData>,
   opts: { strict?: boolean } = {}
-): Schema<TData> {
+): AgenticSchema<TData> {
-  return isSchema(schema) ? schema : createSchemaFromZodSchema(schema, opts)
+  return isAgenticSchema(schema)
+    ? schema
+    : createAgenticSchemaFromZodSchema(schema, opts)
+}
+
+export function asZodOrJsonSchema<TData>(
+  schema: z.Schema<TData> | AgenticSchema<TData>
+): z.Schema<TData> | 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<TData = unknown>(
     safeParse?: types.SafeParseFn<TData>
     source?: any
   } = {}
-): Schema<TData> {
+): AgenticSchema<TData> {
   safeParse ??= (value: unknown) => {
     try {
       const result = parse(value)
@@ -116,10 +130,10 @@ export function createJsonSchema<TData = unknown>(
   }
 }
 
-export function createSchemaFromZodSchema<TData>(
+export function createAgenticSchemaFromZodSchema<TData>(
   zodSchema: z.Schema<TData>,
   opts: { strict?: boolean } = {}
-): Schema<TData> {
+): AgenticSchema<TData> {
   return createJsonSchema(zodToJsonSchema(zodSchema, opts), {
     parse: (value) => {
       return parseStructuredOutput(value, zodSchema)

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<any> | Schema<any>
+export type AIFunctionInputSchema = z.ZodObject<any> | AgenticSchema<any>
 
 // eslint-disable-next-line @typescript-eslint/naming-convention
 export type inferInput<InputSchema extends AIFunctionInputSchema> =
-  InputSchema extends Schema<any>
+  InputSchema extends AgenticSchema<any>
     ? InputSchema['_type']
     : InputSchema extends z.ZodTypeAny
       ? z.infer<InputSchema>

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]: asZodOrJsonSchema(fn.inputSchema),
-        [inputSchemaKey]: isZodSchema(fn.inputSchema)
-          ? fn.inputSchema
-          : asSchema(fn.inputSchema).jsonSchema,
         outputSchema: z.any()
       },
       fn.execute

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

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
     })
   )
 }

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[]