feat(js/ai): added dynamic tools that can be dynamically defined at runtime

js/ai/src/generate.ts

@@ -334,6 +334,8 @@ export async function generate<
   };
   const resolvedFormat = await resolveFormat(registry, resolvedOptions.output);
 
+  registry = maybeRegisterDynamicTools(registry, resolvedOptions);
+
   const params = await toGenerateActionOptions(registry, resolvedOptions);
 
   const tools = await toolsToActionRefs(registry, resolvedOptions.tools);
@@ -362,6 +364,29 @@ export async function generate<
   );
 }
 
+function maybeRegisterDynamicTools<
+  O extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+>(registry: Registry, options: GenerateOptions<O, CustomOptions>): Registry {
+  let hasDynamicTools = false;
+  options?.tools?.forEach((t) => {
+    if (
+      (t as Action).__action &&
+      (t as Action).__action.metadata?.type === 'tool' &&
+      (t as Action).__action.metadata?.dynamic
+    ) {
+      if (!hasDynamicTools) {
+        hasDynamicTools = true;
+        // Create a temporary registry with dynamic tools for the duration of this
+        // generate request.
+        registry = Registry.withParent(registry);
+      }
+      registry.registerAction('tool', t as Action);
+    }
+  });
+  return registry;
+}
+
 export async function toGenerateActionOptions<
   O extends z.ZodTypeAny = z.ZodTypeAny,
   CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,

js/ai/src/generate/resolve-tool-requests.ts

@@ -126,13 +126,18 @@ export async function resolveToolRequest(
 
     return { response };
   } catch (e) {
-    if (e instanceof ToolInterruptError) {
+    if (
+      e instanceof ToolInterruptError ||
+      // There's an inexplicable case when the above type check fails, only in tests.
+      (e as Error).name === 'ToolInterruptError'
+    ) {
+      const ie = e as ToolInterruptError;
       logger.debug(
-        `tool '${toolMap[part.toolRequest?.name].__action.name}' triggered an interrupt${e.metadata ? `: ${JSON.stringify(e.metadata)}` : ''}`
+        `tool '${toolMap[part.toolRequest?.name].__action.name}' triggered an interrupt${ie.metadata ? `: ${JSON.stringify(ie.metadata)}` : ''}`
       );
       const interrupt = {
         toolRequest: part.toolRequest,
-        metadata: { ...part.metadata, interrupt: e.metadata || true },
+        metadata: { ...part.metadata, interrupt: ie.metadata || true },
       };
 
       return { interrupt };

js/ai/src/tool.ts

@@ -15,6 +15,7 @@
  */
 
 import {
+  action,
   Action,
   ActionContext,
   ActionRunOptions,
@@ -24,7 +25,7 @@ import {
   stripUndefinedProps,
   z,
 } from '@genkit-ai/core';
-import { Registry } from '@genkit-ai/core/registry';
+import { HasRegistry, Registry } from '@genkit-ai/core/registry';
 import { parseSchema, toJsonSchema } from '@genkit-ai/core/schema';
 import { setCustomMetadataAttributes } from '@genkit-ai/core/tracing';
 import {
@@ -270,6 +271,15 @@ export function defineTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
       });
     }
   );
+  implementTool(a as ToolAction<I, O>, config, registry);
+  return a as ToolAction<I, O>;
+}
+
+function implementTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
+  a: ToolAction<I, O>,
+  config: ToolConfig<I, O>,
+  registry: Registry
+) {
   (a as ToolAction<I, O>).respond = (interrupt, responseData, options) => {
     assertUnstable(
       registry,
@@ -319,7 +329,6 @@ export function defineTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
       }),
     };
   };
-  return a as ToolAction<I, O>;
 }
 
 /** InterruptConfig defines the options for configuring an interrupt. */
@@ -381,3 +390,35 @@ function interruptTool(registry: Registry) {
     throw new ToolInterruptError(metadata);
   };
 }
+
+/**
+ * Defines a dynamic tool. Dynamic tools are just like regular tools but will not be registered in the
+ * Genkit registry and can be defined dynamically at runtime.
+ */
+export function dynamicTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(
+  ai: HasRegistry,
+  config: ToolConfig<I, O>,
+  fn?: ToolFn<I, O>
+): ToolAction<I, O> {
+  const a = action(
+    ai.registry,
+    {
+      ...config,
+      actionType: 'tool',
+      metadata: { ...(config.metadata || {}), type: 'tool', dynamic: true },
+    },
+    (i, runOptions) => {
+      const interrupt = interruptTool(ai.registry);
+      if (fn) {
+        return fn(i, {
+          ...runOptions,
+          context: { ...runOptions.context },
+          interrupt,
+        });
+      }
+      return interrupt();
+    }
+  );
+  implementTool(a as ToolAction<I, O>, config, ai.registry);
+  return a as ToolAction<I, O>;
+}

js/core/src/registry.ts

@@ -115,21 +115,32 @@ export class Registry {
   private allPluginsInitialized = false;
   public apiStability: 'stable' | 'beta' = 'stable';
 
-  readonly asyncStore = new AsyncStore();
-  readonly dotprompt = new Dotprompt({
-    schemaResolver: async (name) => {
-      const resolvedSchema = await this.lookupSchema(name);
-      if (!resolvedSchema) {
-        throw new GenkitError({
-          message: `Schema '${name}' not found`,
-          status: 'NOT_FOUND',
-        });
-      }
-      return toJsonSchema(resolvedSchema);
-    },
-  });
-
-  constructor(public parent?: Registry) {}
+  readonly asyncStore: AsyncStore;
+  readonly dotprompt: Dotprompt;
+  readonly parent?: Registry;
+
+  constructor(parent?: Registry) {
+    if (parent) {
+      this.parent = parent;
+      this.apiStability = parent?.apiStability;
+      this.asyncStore = parent.asyncStore;
+      this.dotprompt = parent.dotprompt;
+    } else {
+      this.asyncStore = new AsyncStore();
+      this.dotprompt = new Dotprompt({
+        schemaResolver: async (name) => {
+          const resolvedSchema = await this.lookupSchema(name);
+          if (!resolvedSchema) {
+            throw new GenkitError({
+              message: `Schema '${name}' not found`,
+              status: 'NOT_FOUND',
+            });
+          }
+          return toJsonSchema(resolvedSchema);
+        },
+      });
+    }
+  }
 
   /**
    * Creates a new registry overlaid onto the provided registry.

js/genkit/src/common.ts

@@ -108,6 +108,7 @@ export {
   type SessionData,
   type SessionStore,
 } from '@genkit-ai/ai/session';
+export { dynamicTool } from '@genkit-ai/ai/tool';
 export {
   GENKIT_CLIENT_HEADER,
   GENKIT_VERSION,

js/genkit/src/tool.ts

@@ -16,6 +16,7 @@
 
 export {
   asTool,
+  dynamicTool,
   toToolDefinition,
   type ToolAction,
   type ToolArgument,