Added support for output properties on tasks

.changeset/sixty-pandas-deliver.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/sdk": patch
+"@trigger.dev/openai": patch
+---
+
+Adding support for output properties on tasks

apps/webapp/app/presenters/RunPresenter.server.ts

@@ -1,14 +1,11 @@
 import {
-  DisplayPropertiesSchema,
-  DisplayProperty,
-  DisplayPropertySchema,
   ErrorWithStack,
   ErrorWithStackSchema,
   EventSpecificationSchema,
   StyleSchema,
 } from "@trigger.dev/internal";
-import { z } from "zod";
 import { PrismaClient, prisma } from "~/db.server";
+import { mergeProperties } from "~/utils/mergeProperties.server";
 
 type RunOptions = {
   id: string;
@@ -35,6 +32,7 @@ const taskSelect = {
   delayUntil: true,
   description: true,
   properties: true,
+  outputProperties: true,
   error: true,
   startedAt: true,
   completedAt: true,
@@ -67,30 +65,15 @@ export class RunPresenter {
       return undefined;
     }
 
-    //merge the properties from the version and the run, with the run properties taking precedence
-    const mergedProperties = new Map<string, DisplayProperty>();
-    if (run.version.properties) {
-      const properties = DisplayPropertiesSchema.parse(run.version.properties);
-      for (const property of properties) {
-        mergedProperties.set(property.label, property);
-      }
-    }
-    if (run.properties) {
-      const properties = DisplayPropertiesSchema.parse(run.properties);
-      for (const property of properties) {
-        mergedProperties.set(property.label, property);
-      }
-    }
-    if (run.version.eventSpecification) {
-      const eventSpecification = EventSpecificationSchema.parse(
-        run.version.eventSpecification
-      );
-      if (eventSpecification.properties) {
-        for (const properties of eventSpecification.properties) {
-          mergedProperties.set(properties.label, properties);
-        }
-      }
-    }
+    const eventSpecification = EventSpecificationSchema.parse(
+      run.version.eventSpecification
+    );
+
+    const runProperties = mergeProperties(
+      run.version.properties,
+      run.properties,
+      eventSpecification.properties
+    );
 
     const enrichTask = (task: QueryTask) => {
       const { children, ...t } = task;
@@ -98,10 +81,7 @@ export class RunPresenter {
         ...t,
         error: t.error ? ErrorWithStackSchema.parse(t.error) : undefined,
         connection: t.runConnection,
-        properties:
-          t.properties == null
-            ? []
-            : z.array(DisplayPropertySchema).parse(t.properties),
+        properties: mergeProperties(t.properties, t.outputProperties),
         style: t.style ? StyleSchema.parse(t.style) : undefined,
       };
     };
@@ -147,7 +127,7 @@ export class RunPresenter {
       isTest: run.isTest,
       version: run.version.version,
       output: runOutput,
-      properties: Array.from(mergedProperties.values()),
+      properties: runProperties,
       environment: {
         type: run.environment.type,
         slug: run.environment.slug,

apps/webapp/app/presenters/TaskDetailsPresenter.server.ts

@@ -1,5 +1,6 @@
 import { DisplayPropertiesSchema, StyleSchema } from "@trigger.dev/internal";
 import { PrismaClient, prisma } from "~/db.server";
+import { mergeProperties } from "~/utils/mergeProperties.server";
 
 type DetailsProps = {
   id: string;
@@ -56,6 +57,7 @@ export class TaskDetailsPresenter {
         noop: true,
         description: true,
         properties: true,
+        outputProperties: true,
         params: true,
         output: true,
         error: true,
@@ -89,10 +91,7 @@ export class TaskDetailsPresenter {
       ...task,
       connection: task.runConnection,
       params: task.params as Record<string, any>,
-      properties:
-        task.properties == null
-          ? []
-          : DisplayPropertiesSchema.parse(task.properties),
+      properties: mergeProperties(task.properties, task.outputProperties),
       style: task.style ? StyleSchema.parse(task.style) : undefined,
     };
   }

apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.jobs.$jobParam.runs.$runParam.tasks.$taskParam/route.tsx

@@ -3,7 +3,6 @@ import { typedjson, useTypedLoaderData } from "remix-typedjson";
 import { CodeBlock } from "~/components/code/CodeBlock";
 import { DateTime } from "~/components/primitives/DateTime";
 import { Header3 } from "~/components/primitives/Headers";
-import { useLocales } from "~/components/primitives/LocaleProvider";
 import { Paragraph } from "~/components/primitives/Paragraph";
 import {
   Table,
@@ -58,7 +57,6 @@ export const loader = async ({ request, params }: LoaderArgs) => {
 
 export default function Page() {
   const { task } = useTypedLoaderData<typeof loader>();
-  const locales = useLocales();
 
   const {
     name,
@@ -139,7 +137,7 @@ export default function Page() {
         {properties.length > 0 && (
           <div className="mt-4 flex flex-col gap-2">
             <Header3>Properties</Header3>
-            <RunPanelProperties properties={properties} layout="vertical" />
+            <RunPanelProperties properties={properties} layout="horizontal" />
           </div>
         )}
 

apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.$projectParam.jobs.$jobParam.runs.$runParam.trigger/route.tsx

@@ -67,7 +67,7 @@ export default function Page() {
               <Header3>Properties</Header3>
               <RunPanelProperties
                 properties={run.properties}
-                layout="vertical"
+                layout="horizontal"
               />
             </div>
           )}

apps/webapp/app/routes/api.v1.runs.$runId.tasks.$id.complete.ts

@@ -146,6 +146,7 @@ export class CompleteRunTaskService {
             status: "COMPLETED",
             output: taskBody.output ?? undefined,
             completedAt: new Date(),
+            outputProperties: taskBody.properties,
           },
           include: {
             attempts: true,

apps/webapp/app/utils/mergeProperties.server.ts

@@ -0,0 +1,36 @@
+import {
+  DisplayProperty,
+  DisplayPropertiesSchema,
+} from "@trigger.dev/internal";
+
+// Takes a list of potential arrays of DisplayProperties and merges them together so no duplicates exist based on the label
+// For example, if the propertyLists are:
+// [
+//   [{label: "foo", value: "bar"}],
+//   [{label: "foo", value: "baz"}],
+//   [{label: "bar", value: "baz"}],
+// ]
+// The result would be:
+// [
+//   {label: "foo", value: "baz"},
+//   {label: "bar", value: "baz"},
+// ]
+//
+// We will use the DisplayPropertiesSchema zod schema to safely parse the properties and if they aren't valid then we'll just ignore them
+export function mergeProperties(
+  ...propertyLists: Array<unknown>
+): Array<DisplayProperty> {
+  const mergedProperties = new Map<string, DisplayProperty>();
+
+  for (const propertyList of propertyLists) {
+    const properties = DisplayPropertiesSchema.safeParse(propertyList);
+
+    if (properties.success) {
+      for (const property of properties.data) {
+        mergedProperties.set(property.label, property);
+      }
+    }
+  }
+
+  return Array.from(mergedProperties.values());
+}

integrations/openai/src/index.ts

@@ -1,48 +1,8 @@
 import type { IntegrationClient, TriggerIntegration } from "@trigger.dev/sdk";
 import { Configuration, OpenAIApi } from "openai";
-import {
-  backgroundCreateChatCompletion,
-  backgroundCreateCompletion,
-  cancelFineTune,
-  createChatCompletion,
-  createCompletion,
-  createEdit,
-  createEmbedding,
-  createFile,
-  createFineTune,
-  createFineTuneFile,
-  createImage,
-  deleteFineTune,
-  listFiles,
-  listFineTuneEvents,
-  listFineTunes,
-  listModels,
-  retrieveFineTune,
-  retrieveModel,
-} from "./tasks";
+import * as tasks from "./tasks";
 import { OpenAIIntegrationOptions } from "./types";
 
-const tasks = {
-  retrieveModel,
-  listModels,
-  createCompletion,
-  createChatCompletion,
-  backgroundCreateCompletion,
-  backgroundCreateChatCompletion,
-  createEdit,
-  createImage,
-  createEmbedding,
-  createFile,
-  listFiles,
-  createFineTuneFile,
-  createFineTune,
-  listFineTunes,
-  retrieveFineTune,
-  cancelFineTune,
-  listFineTuneEvents,
-  deleteFineTune,
-};
-
 export class OpenAI
   implements TriggerIntegration<IntegrationClient<OpenAIApi, typeof tasks>>
 {

integrations/openai/src/taskUtils.ts

@@ -0,0 +1,34 @@
+import {
+  CreateCompletionResponseUsage,
+  CreateEmbeddingResponseUsage,
+} from "openai";
+
+export function createTaskUsageProperties(
+  usage:
+    | CreateCompletionResponseUsage
+    | CreateEmbeddingResponseUsage
+    | undefined
+) {
+  if (!usage) {
+    return;
+  }
+
+  return [
+    {
+      label: "Prompt Usage",
+      text: String(usage.prompt_tokens),
+    },
+    ...("completion_tokens" in usage
+      ? [
+          {
+            label: "Completion Usage",
+            text: String(usage.completion_tokens),
+          },
+        ]
+      : []),
+    {
+      label: "Total Usage",
+      text: String(usage.total_tokens),
+    },
+  ];
+}

integrations/openai/src/tasks.ts

@@ -15,6 +15,7 @@ import {
   fileFromString,
   truncate,
 } from "@trigger.dev/integration-kit";
+import { createTaskUsageProperties } from "./taskUtils";
 
 type OpenAIClientType = InstanceType<typeof OpenAIApi>;
 
@@ -76,8 +77,12 @@ export const createCompletion: AuthenticatedTask<
   Prettify<CreateCompletionRequest>,
   Prettify<Awaited<ReturnType<OpenAIClientType["createCompletion"]>>["data"]>
 > = {
-  run: async (params, client) => {
-    return client.createCompletion(params).then((res) => res.data);
+  run: async (params, client, task) => {
+    const response = await client.createCompletion(params);
+
+    task.outputProperties = createTaskUsageProperties(response.data.usage);
+
+    return response.data;
   },
   init: (params) => {
     return {
@@ -105,7 +110,7 @@ export const backgroundCreateCompletion: AuthenticatedTask<
   OpenAIIntegrationAuth
 > = {
   run: async (params, client, task, io, auth) => {
-    return io.backgroundFetch<CreateCompletionResponseData>(
+    const response = await io.backgroundFetch<CreateCompletionResponseData>(
       "background",
       "https://api.openai.com/v1/completions",
       {
@@ -120,6 +125,10 @@ export const backgroundCreateCompletion: AuthenticatedTask<
         body: JSON.stringify(params),
       }
     );
+
+    task.outputProperties = createTaskUsageProperties(response.usage);
+
+    return response;
   },
   init: (params) => {
     return {
@@ -145,8 +154,12 @@ export const createChatCompletion: AuthenticatedTask<
   Prettify<CreateChatCompletionRequest>,
   Prettify<CreateChatCompetionResponseData>
 > = {
-  run: async (params, client) => {
-    return client.createChatCompletion(params).then((res) => res.data);
+  run: async (params, client, task) => {
+    const response = await client.createChatCompletion(params);
+
+    task.outputProperties = createTaskUsageProperties(response.data.usage);
+
+    return response.data;
   },
   init: (params) => {
     return {
@@ -170,7 +183,7 @@ export const backgroundCreateChatCompletion: AuthenticatedTask<
   OpenAIIntegrationAuth
 > = {
   run: async (params, client, task, io, auth) => {
-    return io.backgroundFetch<CreateChatCompetionResponseData>(
+    const response = await io.backgroundFetch<CreateChatCompetionResponseData>(
       "background",
       "https://api.openai.com/v1/chat/completions",
       {
@@ -203,6 +216,10 @@ export const backgroundCreateChatCompletion: AuthenticatedTask<
         },
       }
     );
+
+    task.outputProperties = createTaskUsageProperties(response.usage);
+
+    return response;
   },
   init: (params) => {
     return {
@@ -228,8 +245,12 @@ export const createEdit: AuthenticatedTask<
   Prettify<CreateEditRequest>,
   CreateEditResponseData
 > = {
-  run: async (params, client) => {
-    return client.createEdit(params).then((res) => res.data);
+  run: async (params, client, task) => {
+    const response = await client.createEdit(params);
+
+    task.outputProperties = createTaskUsageProperties(response.data.usage);
+
+    return response.data;
   },
   init: (params) => {
     let properties = [
@@ -270,7 +291,9 @@ export const createImage: AuthenticatedTask<
   CreateImageResponseData
 > = {
   run: async (params, client, task) => {
-    return client.createImage(params).then((res) => res.data);
+    const response = await client.createImage(params);
+
+    return response.data;
   },
   init: (params) => {
     let properties = [
@@ -319,8 +342,12 @@ export const createEmbedding: AuthenticatedTask<
   Prettify<CreateEmbeddingRequest>,
   CreateEmbeddingResponseData
 > = {
-  run: async (params, client) => {
-    return client.createEmbedding(params).then((res) => res.data);
+  run: async (params, client, task) => {
+    const response = await client.createEmbedding(params);
+
+    task.outputProperties = createTaskUsageProperties(response.data.usage);
+
+    return response.data;
   },
   init: (params) => {
     return {

packages/database/prisma/migrations/20230708193746_add_output_properties_to_tasks/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "Task" ADD COLUMN     "outputProperties" JSONB;

packages/database/prisma/schema.prisma

@@ -760,14 +760,15 @@ model Task {
   delayUntil DateTime?
   noop       Boolean    @default(false)
 
-  description String?
-  properties  Json?
-  params      Json?
-  output      Json?
-  error       String?
-  redact      Json?
-  style       Json?
-  operation   String?
+  description      String?
+  properties       Json?
+  outputProperties Json?
+  params           Json?
+  output           Json?
+  error            String?
+  redact           Json?
+  style            Json?
+  operation        String?
 
   startedAt   DateTime?
   completedAt DateTime?

packages/internal/src/schemas/api.ts

@@ -13,6 +13,7 @@ import {
 } from "./schedules";
 import { CachedTaskSchema, ServerTaskSchema, TaskSchema } from "./tasks";
 import { EventSpecificationSchema, TriggerMetadataSchema } from "./triggers";
+import { Prettify } from "../types";
 
 export const UpdateTriggerSourceBodySchema = z.object({
   registeredEvents: z.array(z.string()),
@@ -530,7 +531,9 @@ export const CompleteTaskBodyInputSchema = RunTaskBodyInputSchema.pick({
   ),
 });
 
-export type CompleteTaskBodyInput = z.input<typeof CompleteTaskBodyInputSchema>;
+export type CompleteTaskBodyInput = Prettify<
+  z.input<typeof CompleteTaskBodyInputSchema>
+>;
 export type CompleteTaskBodyOutput = z.infer<
   typeof CompleteTaskBodyInputSchema
 >;

packages/internal/src/schemas/tasks.ts

@@ -23,6 +23,7 @@ export const TaskSchema = z.object({
   status: TaskStatusSchema,
   description: z.string().optional().nullable(),
   properties: z.array(DisplayPropertySchema).optional().nullable(),
+  outputProperties: z.array(DisplayPropertySchema).optional().nullable(),
   params: DeserializedJsonSchema.optional().nullable(),
   output: DeserializedJsonSchema.optional().nullable(),
   error: z.string().optional().nullable(),

packages/trigger-sdk/src/io.ts

@@ -558,6 +558,7 @@ export class IO {
 
         await this._apiClient.completeTask(this._id, task.id, {
           output: result ?? undefined,
+          properties: task.outputProperties ?? undefined,
         });
 
         return result;