feat: tool execution state

packages/core/src/defineTool.ts

@@ -1,31 +1,25 @@
-import type { EmptyObject } from '@nzyme/types';
-
 import * as s from '@agentscript-ai/schema';
 
 const TOOL_SYMBOL = Symbol('tool');
 
-type ToolInputOptions = s.ObjectSchemaProps | s.NonNullish<s.ObjectSchema>;
+type ToolInputSchema = s.NonNullish<s.ObjectSchema> | s.VoidSchema;
+type ToolInputOptions = s.ObjectSchemaProps | s.NonNullish<s.ObjectSchema> | undefined;
 
-type ToolInputSchema<TIn extends ToolInputOptions> =
+type ToolInputSchemaFromOptions<TIn extends ToolInputOptions = ToolInputOptions> =
     TIn extends s.NonNullish<s.ObjectSchema>
         ? TIn
         : TIn extends s.ObjectSchemaProps
           ? s.ObjectSchema<{ props: TIn; nullable: false; optional: false }>
-          : never;
-
-type ToolInputValue<TIn extends ToolInputOptions> = TIn extends s.ObjectSchema
-    ? s.Infer<TIn>
-    : TIn extends s.ObjectSchemaProps
-      ? s.ObjectSchemaProps extends TIn
-          ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            any
-          : s.ObjectSchemaPropsValue<TIn>
-      : never;
+          : s.VoidSchema;
 
 /**
  * Options for {@link defineTool}.
  */
-export type ToolOptions<TIn extends ToolInputOptions, TOut extends s.Schema> = {
+export type ToolOptions<
+    TInput extends ToolInputOptions,
+    TOutput extends s.Schema,
+    TState extends ToolInputOptions,
+> = {
     /**
      * Description of the tool.
      * Should be descriptive and concise, so that LLM can understand what the tool does.
@@ -34,24 +28,33 @@ export type ToolOptions<TIn extends ToolInputOptions, TOut extends s.Schema> = {
     /**
      * Arguments for the tool.
      */
-    input?: TIn;
+    input?: TInput;
     /**
      * Schema of the return value of the tool.
      * @default s.void()
      */
-    output?: TOut;
+    output?: TOutput;
+    /**
+     * Schema of the state of the tool.
+     */
+    state?: TState;
     /**
      * Handler for the tool.
      */
-    handler: ToolHandler<TIn, TOut>;
+    handler: ToolHandler<
+        s.Infer<ToolInputSchemaFromOptions<TInput>>,
+        s.Infer<TOutput>,
+        s.Infer<ToolInputSchemaFromOptions<TState>>
+    >;
 };
 
 /**
  * Tool definition.
  */
 export type ToolDefinition<
-    TIn extends s.NonNullish<s.ObjectSchema> = s.NonNullish<s.ObjectSchemaAny>,
-    TOut extends s.Schema = s.SchemaAny,
+    TInput extends ToolInputSchema = ToolInputSchema,
+    TOutput extends s.Schema = s.SchemaAny,
+    TState extends ToolInputSchema = ToolInputSchema,
 > = {
     /**
      * Description of the tool.
@@ -61,19 +64,23 @@ export type ToolDefinition<
     /**
      * Arguments for the tool.
      */
-    input: TIn;
+    input: TInput;
     /**
      * Whether the tool has a single argument.
      */
     singleArg: boolean;
     /**
      * Schema of the return value of the tool.
      */
-    output: TOut;
+    output: TOutput;
+    /**
+     * Schema of the state of the tool.
+     */
+    state: TState;
     /**
      * Handler for the tool.
      */
-    handler: ToolHandler<TIn, TOut>;
+    handler: ToolHandler<s.Infer<TInput>, s.Infer<TOutput>, s.Infer<TState>>;
     /**
      * Symbol to indicate that the value is a tool.
      * @internal
@@ -84,54 +91,76 @@ export type ToolDefinition<
 /**
  * Parameters for the tool handler.
  */
-export type ToolHandlerParams<TIn extends ToolInputOptions> = {
+export type ToolContext<TInput, TState> = {
     /**
      * Resolved arguments for the tool.
+     * First you need to define the input schema in {@link defineTool} options.
      */
-    input: ToolInputValue<TIn>;
+    input: TInput;
+    /**
+     * State of the tool.
+     */
+    state: TState;
 };
 
 /**
  * Handler for the tool.
  */
-export type ToolHandler<TIn extends ToolInputOptions, TOut extends s.Schema> = (
-    params: ToolHandlerParams<TIn>,
-) => s.Infer<TOut> | Promise<s.Infer<TOut>>;
+export type ToolHandler<TInput, TOutput, TState> = (
+    ctx: ToolContext<TInput, TState>,
+) => TOutput | Promise<TOutput>;
 
 /**
  * Define a tool.
  * @param options - Options for the tool.
  * @returns Defined tool.
  */
 export function defineTool<
-    TIn extends ToolInputOptions = EmptyObject,
-    TOut extends s.Schema = s.Schema<void>,
->(options: ToolOptions<TIn, TOut>): ToolDefinition<ToolInputSchema<TIn>, TOut> {
-    let input: ToolInputSchema<TIn>;
+    TInput extends ToolInputOptions = undefined,
+    TOutput extends s.Schema = s.VoidSchema,
+    TState extends ToolInputOptions = undefined,
+>(
+    options: ToolOptions<TInput, TOutput, TState>,
+): ToolDefinition<ToolInputSchemaFromOptions<TInput>, TOutput, ToolInputSchemaFromOptions<TState>> {
+    let input: ToolInputSchemaFromOptions<TInput>;
     let singleArg = false;
 
     if (!options.input) {
-        input = s.object({ props: {} }) as ToolInputSchema<TIn>;
+        input = s.void() as ToolInputSchemaFromOptions<TInput>;
     } else if (s.isSchema(options.input, s.object)) {
         if (options.input.nullable || options.input.optional) {
             throw new Error('Input schema must not be nullable or optional');
         }
 
-        input = options.input as ToolInputSchema<TIn>;
+        input = options.input as ToolInputSchemaFromOptions<TInput>;
         // if a full schema is provided, we assume it's a single argument
         singleArg = true;
     } else {
-        input = s.object({ props: options.input }) as ToolInputSchema<TIn>;
+        input = s.object({ props: options.input }) as ToolInputSchemaFromOptions<TInput>;
         // make the tool singleArg if there are more than 2 arguments
         singleArg = Object.keys(options.input).length > 2;
     }
 
+    let state: ToolInputSchemaFromOptions<TState>;
+    if (!options.state) {
+        state = s.void() as ToolInputSchemaFromOptions<TState>;
+    } else if (s.isSchema(options.state, s.object)) {
+        if (options.state.nullable || options.state.optional) {
+            throw new Error('State schema must not be nullable or optional');
+        }
+
+        state = options.state as ToolInputSchemaFromOptions<TState>;
+    } else {
+        state = s.object({ props: options.state }) as ToolInputSchemaFromOptions<TState>;
+    }
+
     return {
         description: options.description,
         input,
         singleArg,
-        output: options.output ?? (s.void() as TOut),
-        handler: options.handler as ToolHandler<ToolInputSchema<TIn>, TOut>,
+        output: options.output ?? (s.void() as TOutput),
+        state,
+        handler: options.handler,
         [TOOL_SYMBOL]: true,
     };
 }

packages/core/src/modules/renderTool.ts

@@ -1,3 +1,4 @@
+import * as s from '@agentscript-ai/schema';
 import { normalizeText } from '@agentscript-ai/utils';
 
 import type { ToolDefinition } from '../defineTool.js';
@@ -23,29 +24,31 @@ export function renderTool(options: RenderToolOptions) {
 
     const description = normalizeText(tool.description);
 
-    if (tool.singleArg) {
-        const input = tool.input;
-        const inputTypeName = renderType({
-            schema: input,
-            ctx,
-            nameHint: `${name}Params`,
-        });
+    const input = tool.input;
+    if (s.isSchema(input, s.object)) {
+        if (tool.singleArg) {
+            const inputTypeName = renderType({
+                schema: input,
+                ctx,
+                nameHint: `${name}Params`,
+            });
 
-        if (input.description) {
-            description.push(renderDocDirective(`param params -`, input.description));
-        }
-
-        args = `params: ${inputTypeName}`;
-    } else {
-        for (const [name, arg] of Object.entries(tool.input.props)) {
-            if (args.length > 0) {
-                args += ', ';
+            if (input.description) {
+                description.push(renderDocDirective(`param params -`, input.description));
             }
 
-            args += `${name}: ${renderType({ schema: arg, ctx })}`;
+            args = `params: ${inputTypeName}`;
+        } else {
+            for (const [name, arg] of Object.entries(input.props)) {
+                if (args.length > 0) {
+                    args += ', ';
+                }
+
+                args += `${name}: ${renderType({ schema: arg, ctx })}`;
 
-            if (arg.description) {
-                description.push(renderDocDirective(`param ${name} -`, arg.description));
+                if (arg.description) {
+                    description.push(renderDocDirective(`param ${name} -`, arg.description));
+                }
             }
         }
     }

packages/core/src/runtime/executeAgent.ts

@@ -1,5 +1,6 @@
 import type { Constructor } from '@nzyme/types';
 
+import * as s from '@agentscript-ai/schema';
 import { validateOrThrow } from '@agentscript-ai/schema';
 
 import { RuntimeError } from './RuntimeError.js';
@@ -410,24 +411,33 @@ async function runFunctionCustom(
     }
 
     let argObject: Record<string, unknown>;
-
     if (func.singleArg) {
         argObject = args[0] as Record<string, unknown>;
     } else {
         argObject = {};
-        const argProps = Object.entries(func.input.props);
 
-        for (let i = 0; i < argProps.length; i++) {
-            const arg = args[i];
-            const argName = argProps[i][0];
+        if (s.isSchema(func.input, s.object)) {
+            const argProps = Object.entries(func.input.props);
+
+            for (let i = 0; i < argProps.length; i++) {
+                const arg = args[i];
+                const argName = argProps[i][0];
 
-            argObject[argName] = arg;
+                argObject[argName] = arg;
+            }
         }
     }
 
     validateOrThrow(func.input, argObject);
 
-    const result: unknown = func.handler({ input: argObject });
+    // Prepare state for the tool execution
+    let state = frame.state as Record<string, unknown> | undefined;
+    if (!state || !s.isSchema(state, s.object)) {
+        state = s.coerce(func.state) as Record<string, unknown>;
+        frame.state = state;
+    }
+
+    const result: unknown = func.handler({ input: argObject, state });
     if (result instanceof Promise) {
         frame.value = await result;
         controller.tick();

packages/core/src/runtime/runtimeTypes.ts

@@ -1,23 +1,65 @@
+/**
+ * Stack frame.
+ */
 export interface StackFrame {
+    /**
+     * Started at.
+     */
     startedAt: number;
+    /**
+     * Completed at.
+     */
     completedAt?: number;
+    /**
+     * Variables.
+     */
     variables?: Record<string, unknown>;
+    /**
+     * Parent frame.
+     */
     parent?: StackFrame;
+    /**
+     * Error, if any.
+     */
     error?: string;
+    /**
+     * Value, ie the result of the function.
+     */
     value?: unknown;
+    /**
+     * State, ie the state of the tool.
+     */
+    state?: unknown;
+    /**
+     * Children frames.
+     */
     children?: StackFrame[];
 }
 
+/**
+ * Stack block frame.
+ */
 export interface StackBlockFrame extends StackFrame {
+    /**
+     * Variables.
+     */
     variables: Record<string, unknown>;
+    /**
+     * Frames.
+     */
     frames: StackFrame[];
 }
 
-export interface StackFunctionFrame extends StackFrame {
-    state: unknown;
-}
-
+/**
+ * Stack loop frame.
+ */
 export interface StackLoopFrame extends StackFrame {
+    /**
+     * Item name.
+     */
     itemName: string;
+    /**
+     * Item blocks.
+     */
     itemBlocks: StackBlockFrame[];
 }