Added AI tool tasks, descriptions to tasks

Author: ericallam

apps/webapp/app/v3/services/createBackgroundWorker.server.ts

@@ -150,6 +150,7 @@ export async function createBackgroundTasks(
           runtimeEnvironmentId: worker.runtimeEnvironmentId,
           workerId: worker.id,
           slug: task.id,
+          description: task.description,
           filePath: task.filePath,
           exportName: task.exportName,
           retryConfig: task.retry,

internal-packages/database/prisma/migrations/20241112111156_add_description_to_background_worker_task/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "BackgroundWorkerTask" ADD COLUMN     "description" TEXT;

internal-packages/database/prisma/schema.prisma

@@ -1610,6 +1610,8 @@ model BackgroundWorkerTask {
   id   String @id @default(cuid())
   slug String
 
+  description String?
+
   friendlyId String @unique
 
   filePath   String

packages/core/package.json

@@ -204,16 +204,18 @@
     "nanoid": "^3.3.4",
     "socket.io-client": "4.7.5",
     "superjson": "^2.2.1",
-    "zod": "3.22.3",
     "zod-error": "1.5.0",
-    "zod-validation-error": "^1.5.0"
+    "zod-validation-error": "^1.5.0",
+    "zod": "3.22.3"
   },
   "devDependencies": {
+    "@ai-sdk/provider-utils": "^1.0.22",
     "@arethetypeswrong/cli": "^0.15.4",
     "@epic-web/test-server": "^0.1.0",
     "@types/humanize-duration": "^3.27.1",
     "@types/node": "20.14.14",
     "@types/readable-stream": "^4.0.14",
+    "ai": "^3.4.33",
     "defu": "^6.1.4",
     "esbuild": "^0.23.0",
     "rimraf": "^3.0.2",

packages/core/src/v3/schemas/resources.ts

@@ -4,6 +4,7 @@ import { MachineConfig } from "./common.js";
 
 export const TaskResource = z.object({
   id: z.string(),
+  description: z.string().optional(),
   filePath: z.string(),
   exportName: z.string(),
   queue: QueueOptions.optional(),

packages/core/src/v3/schemas/schemas.ts

@@ -149,6 +149,7 @@ export const ScheduleMetadata = z.object({
 
 const taskMetadata = {
   id: z.string(),
+  description: z.string().optional(),
   queue: QueueOptions.optional(),
   retry: RetryOptions.optional(),
   machine: MachineConfig.optional(),

packages/core/src/v3/types/index.ts

@@ -5,6 +5,7 @@ import { Prettify } from "./utils.js";
 export * from "./utils.js";
 export * from "./tasks.js";
 export * from "./idempotencyKeys.js";
+export * from "./tools.js";
 
 type ResolveEnvironmentVariablesOptions = {
   variables: Record<string, string> | Array<{ name: string; value: string }>;

packages/core/src/v3/types/tasks.ts

@@ -1,17 +1,20 @@
+import type { Schema as AISchema } from "ai";
+import { z } from "zod";
 import { SerializableJson } from "../../schemas/json.js";
+import { TriggerApiRequestOptions } from "../apiClient/index.js";
 import { RunTags } from "../schemas/api.js";
-import { QueueOptions } from "../schemas/schemas.js";
-import { IdempotencyKey } from "./idempotencyKeys.js";
 import {
   MachineCpu,
   MachineMemory,
   RetryOptions,
   TaskMetadata,
   TaskRunContext,
 } from "../schemas/index.js";
+import { QueueOptions } from "../schemas/schemas.js";
+import { IdempotencyKey } from "./idempotencyKeys.js";
+import { AnySchemaParseFn, inferSchemaIn, inferSchemaOut, Schema } from "./schemas.js";
 import { Prettify } from "./utils.js";
-import { AnySchemaParseFn, inferSchemaOut, Schema } from "./schemas.js";
-import { TriggerApiRequestOptions } from "../apiClient/index.js";
+import { inferToolParameters, ToolTaskParameters } from "./tools.js";
 
 type RequireOne<T, K extends keyof T> = {
   [X in Exclude<keyof T, K>]?: T[X];
@@ -150,6 +153,8 @@ type CommonTaskOptions<
   /** An id for your task. This must be unique inside your project and not change between versions.  */
   id: TIdentifier;
 
+  description?: string;
+
   /** The retry settings when an uncaught error is thrown.
        *
        * If omitted it will use the values in your `trigger.config.ts` file.
@@ -337,6 +342,15 @@ export type TaskWithSchemaOptions<
   schema?: TSchema;
 };
 
+export type TaskWithToolOptions<
+  TIdentifier extends string,
+  TParameters extends ToolTaskParameters,
+  TOutput = unknown,
+  TInitOutput extends InitOutput = any,
+> = CommonTaskOptions<TIdentifier, inferToolParameters<TParameters>, TOutput, TInitOutput> & {
+  parameters: TParameters;
+};
+
 declare const __output: unique symbol;
 declare const __payload: unique symbol;
 type BrandRun<P, O> = { [__output]: O; [__payload]: P };
@@ -413,6 +427,9 @@ export interface Task<TIdentifier extends string, TInput = void, TOutput = any>
    * The id of the task.
    */
   id: TIdentifier;
+
+  description?: string;
+
   /**
    * Trigger a task with the given payload, and continue without waiting for the result. If you want to wait for the result, use `triggerAndWait`. Returns the id of the triggered task run.
    * @param payload
@@ -479,6 +496,26 @@ export interface Task<TIdentifier extends string, TInput = void, TOutput = any>
   batchTriggerAndWait: (items: Array<BatchItem<TInput>>) => Promise<BatchResult<TOutput>>;
 }
 
+export interface TaskWithSchema<
+  TIdentifier extends string,
+  TSchema extends TaskSchema | undefined = undefined,
+  TOutput = any,
+> extends Task<TIdentifier, inferSchemaIn<TSchema>, TOutput> {
+  schema?: TSchema;
+}
+
+export interface ToolTask<
+  TIdentifier extends string,
+  TParameters extends ToolTaskParameters,
+  TOutput = any,
+> extends Task<TIdentifier, inferToolParameters<TParameters>, TOutput> {
+  tool: {
+    parameters: TParameters;
+    description?: string;
+    execute: (args: inferToolParameters<TParameters>) => Promise<TOutput>;
+  };
+}
+
 export type AnyTask = Task<string, any, any>;
 
 export type TaskPayload<TTask extends AnyTask> = TTask extends Task<string, infer TInput, any>

packages/core/src/v3/types/tools.ts

@@ -0,0 +1,36 @@
+import { z } from "zod";
+import type { Schema as AISchema } from "ai";
+import { Schema } from "./schemas.js";
+
+export type ToolTaskParameters = z.ZodTypeAny | AISchema<any>;
+
+export type inferToolParameters<PARAMETERS extends ToolTaskParameters> =
+  PARAMETERS extends AISchema<any>
+    ? PARAMETERS["_type"]
+    : PARAMETERS extends z.ZodTypeAny
+    ? z.infer<PARAMETERS>
+    : never;
+
+export function convertToolParametersToSchema<TToolParameters extends ToolTaskParameters>(
+  toolParameters: TToolParameters
+): Schema {
+  return toolParameters instanceof z.ZodSchema
+    ? toolParameters
+    : convertAISchemaToTaskSchema(toolParameters);
+}
+
+function convertAISchemaToTaskSchema(schema: AISchema<any>): Schema {
+  return (payload: unknown) => {
+    const result = schema.validate?.(payload);
+
+    if (!result) {
+      throw new Error("Invalid payload");
+    }
+
+    if (!result.success) {
+      throw result.error;
+    }
+
+    return result.value;
+  };
+}

packages/trigger-sdk/package.json

@@ -57,8 +57,7 @@
     "terminal-link": "^3.0.0",
     "ulid": "^2.3.0",
     "uuid": "^9.0.0",
-    "ws": "^8.11.0",
-    "zod": "3.22.3"
+    "ws": "^8.11.0"
   },
   "devDependencies": {
     "@arethetypeswrong/cli": "^0.15.4",
@@ -67,12 +66,17 @@
     "@types/slug": "^5.0.3",
     "@types/uuid": "^9.0.0",
     "@types/ws": "^8.5.3",
+    "ai": "^3.4.33",
     "encoding": "^0.1.13",
     "rimraf": "^3.0.2",
     "tshy": "^3.0.2",
     "tsx": "4.17.0",
     "typed-emitter": "^2.1.0",
-    "typescript": "^5.5.4"
+    "typescript": "^5.5.4",
+    "zod": "3.22.3"
+  },
+  "peerDependencies": {
+    "zod": "^3.0.0"
   },
   "engines": {
     "node": ">=18.20.0"

packages/trigger-sdk/src/v3/shared.ts

@@ -11,6 +11,7 @@ import {
   ApiRequestOptions,
   BatchTaskRunExecutionResult,
   conditionallyImportPacket,
+  convertToolParametersToSchema,
   createErrorTaskError,
   defaultRetryOptions,
   getSchemaParseFn,
@@ -32,6 +33,7 @@ import {
 import { IdempotencyKey, idempotencyKeys, isIdempotencyKey } from "./idempotencyKeys.js";
 import { PollOptions, runs } from "./runs.js";
 import { tracer } from "./tracer.js";
+import type { Schema as AISchema, CoreTool } from "ai";
 
 import type {
   AnyRunHandle,
@@ -60,9 +62,14 @@ import type {
   TaskRunOptions,
   TaskRunResult,
   TaskSchema,
+  TaskWithSchema,
   TaskWithSchemaOptions,
+  TaskWithToolOptions,
+  ToolTask,
+  ToolTaskParameters,
   TriggerApiRequestOptions,
 } from "@trigger.dev/core/v3";
+import { z } from "zod";
 
 export type {
   AnyRunHandle,
@@ -111,6 +118,7 @@ export function createTask<
 
   const task: Task<TIdentifier, TInput, TOutput> = {
     id: params.id,
+    description: params.description,
     trigger: async (payload, options) => {
       const taskMetadata = taskCatalog.getTaskManifest(params.id);
 
@@ -183,6 +191,7 @@ export function createTask<
 
   taskCatalog.registerTaskMetadata({
     id: params.id,
+    description: params.description,
     queue: params.queue,
     retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined,
     machine: params.machine,
@@ -205,14 +214,39 @@ export function createTask<
   return task;
 }
 
+export function createToolTask<
+  TIdentifier extends string,
+  TParameters extends ToolTaskParameters,
+  TOutput = unknown,
+  TInitOutput extends InitOutput = any,
+>(
+  params: TaskWithToolOptions<TIdentifier, TParameters, TOutput, TInitOutput>
+): ToolTask<TIdentifier, TParameters, TOutput> {
+  const task = createSchemaTask({
+    ...params,
+    schema: convertToolParametersToSchema(params.parameters),
+  });
+
+  return {
+    ...task,
+    tool: {
+      parameters: params.parameters,
+      description: params.description,
+      execute: async (args: any) => {
+        return task.triggerAndWait(args).unwrap();
+      },
+    },
+  };
+}
+
 export function createSchemaTask<
   TIdentifier extends string,
   TSchema extends TaskSchema | undefined = undefined,
   TOutput = unknown,
   TInitOutput extends InitOutput = any,
 >(
   params: TaskWithSchemaOptions<TIdentifier, TSchema, TOutput, TInitOutput>
-): Task<TIdentifier, inferSchemaIn<TSchema>, TOutput> {
+): TaskWithSchema<TIdentifier, TSchema, TOutput> {
   const customQueue = params.queue
     ? queue({
         name: params.queue?.name ?? `task/${params.id}`,
@@ -224,8 +258,10 @@ export function createSchemaTask<
     ? getSchemaParseFn<inferSchemaIn<TSchema>>(params.schema)
     : undefined;
 
-  const task: Task<TIdentifier, inferSchemaIn<TSchema>, TOutput> = {
+  const task: TaskWithSchema<TIdentifier, TSchema, TOutput> = {
     id: params.id,
+    description: params.description,
+    schema: params.schema,
     trigger: async (payload, options, requestOptions) => {
       const taskMetadata = taskCatalog.getTaskManifest(params.id);
 
@@ -299,6 +335,7 @@ export function createSchemaTask<
 
   taskCatalog.registerTaskMetadata({
     id: params.id,
+    description: params.description,
     queue: params.queue,
     retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined,
     machine: params.machine,

packages/trigger-sdk/src/v3/tasks.ts

@@ -3,6 +3,7 @@ import {
   batchTriggerAndWait,
   createTask,
   createSchemaTask,
+  createToolTask,
   SubtaskUnwrapError,
   trigger,
   triggerAndPoll,
@@ -65,6 +66,8 @@ export const task = createTask;
 
 export const schemaTask = createSchemaTask;
 
+export const toolTask = createToolTask;
+
 export const tasks = {
   trigger,
   triggerAndPoll,