Add custom baseURL and idempotency key support to OpenAI integration (#682)

This commit introduces the ability to set custom base URLs, which allows for the use of different APIs such as Perplexity. It also ensures that any OpenAI-related task now includes an idempotency key when making requests. Works for background tasks as well.

.changeset/rude-kids-joke.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/openai": patch
+---
+
+Add ability to set baseURL to use different APIs (like Perplexity)

integrations/openai/src/chat.ts

@@ -14,7 +14,9 @@ export class Chat {
       return this.runTask(
         key,
         async (client, task) => {
-          const response = await client.chat.completions.create(params);
+          const response = await client.chat.completions.create(params, {
+            idempotencyKey: task.idempotencyKey,
+          });
           task.outputProperties = createTaskUsageProperties(response.usage);
           return response;
         },
@@ -38,15 +40,23 @@ export class Chat {
       return this.runTask(
         key,
         async (client, task, io) => {
+          let baseURL = client.baseURL ?? "https://api.openai.com/v1";
+          if (baseURL.endsWith("/")) {
+            baseURL = baseURL.slice(0, -1);
+          }
+
+          const chatCompletionsURL = `${baseURL}/chat/completions`;
+
           const response = await io.backgroundFetch<OpenAI.Chat.ChatCompletion>(
             "background",
-            "https://api.openai.com/v1/chat/completions",
+            chatCompletionsURL,
             {
               method: "POST",
               headers: {
                 "Content-Type": "application/json",
                 Authorization: redactString`Bearer ${client.apiKey}`,
                 ...(client.organization ? { "OpenAI-Organization": client.organization } : {}),
+                "Idempotency-Key": task.idempotencyKey,
               },
               body: JSON.stringify(params),
             },

integrations/openai/src/completions.ts

@@ -17,7 +17,9 @@ export class Completions {
     return this.runTask(
       key,
       async (client, task) => {
-        const response = await client.completions.create(params);
+        const response = await client.completions.create(params, {
+          idempotencyKey: task.idempotencyKey,
+        });
         task.outputProperties = createTaskUsageProperties(response.usage);
         return response;
       },
@@ -41,15 +43,23 @@ export class Completions {
     return this.runTask(
       key,
       async (client, task, io) => {
+        let baseURL = client.baseURL ?? "https://api.openai.com/v1";
+        if (baseURL.endsWith("/")) {
+          baseURL = baseURL.slice(0, -1);
+        }
+
+        const chatCompletionsURL = `${baseURL}/chat/completions`;
+
         const response = await io.backgroundFetch<OpenAI.Completion>(
           "background",
-          "https://api.openai.com/v1/completions",
+          chatCompletionsURL,
           {
             method: "POST",
             headers: {
               "Content-Type": "application/json",
               Authorization: redactString`Bearer ${client.apiKey}`,
               ...(client.organization ? { "OpenAI-Organization": client.organization } : {}),
+              "Idempotency-Key": task.idempotencyKey,
             },
             body: JSON.stringify(params),
           },

integrations/openai/src/edits.ts

@@ -37,7 +37,7 @@ export class Edits {
     return this.runTask(
       key,
       async (client, task) => {
-        const response = await client.edits.create(params);
+        const response = await client.edits.create(params, { idempotencyKey: task.idempotencyKey });
         task.outputProperties = createTaskUsageProperties(response.usage);
         return response;
       },

integrations/openai/src/embeddings.ts

@@ -17,7 +17,9 @@ export class Embeddings {
     return this.runTask(
       key,
       async (client, task) => {
-        const response = await client.embeddings.create(params);
+        const response = await client.embeddings.create(params, {
+          idempotencyKey: task.idempotencyKey,
+        });
         task.outputProperties = createTaskUsageProperties(response.usage);
         return response;
       },

integrations/openai/src/files.ts

@@ -36,7 +36,10 @@ export class Files {
           file = params.file;
         }
 
-        return client.files.create({ file, purpose: params.purpose });
+        return client.files.create(
+          { file, purpose: params.purpose },
+          { idempotencyKey: task.idempotencyKey }
+        );
       },
       {
         name: "Create file",
@@ -81,7 +84,10 @@ export class Files {
           params.fileName
         );
 
-        return client.files.create({ file, purpose: "fine-tune" });
+        return client.files.create(
+          { file, purpose: "fine-tune" },
+          { idempotencyKey: task.idempotencyKey }
+        );
       },
       {
         name: "Create fine tune file",

integrations/openai/src/fineTunes.ts

@@ -41,7 +41,7 @@ export class FineTunes {
     return this.runTask(
       key,
       async (client, task) => {
-        return client.fineTunes.create(params);
+        return client.fineTunes.create(params, { idempotencyKey: task.idempotencyKey });
       },
       {
         name: "Create fine tune",
@@ -94,7 +94,7 @@ export class FineTunes {
     return this.runTask(
       key,
       async (client, task) => {
-        return client.fineTunes.cancel(params.fineTuneId);
+        return client.fineTunes.cancel(params.fineTuneId, { idempotencyKey: task.idempotencyKey });
       },
       {
         name: "Cancel fine tune",
@@ -168,7 +168,7 @@ export class FineTunes {
       return this.runTask(
         key,
         async (client, task) => {
-          return client.fineTuning.jobs.create(params);
+          return client.fineTuning.jobs.create(params, { idempotencyKey: task.idempotencyKey });
         },
         {
           name: "Create Fine Tuning Job",
@@ -207,7 +207,7 @@ export class FineTunes {
       return this.runTask(
         key,
         async (client, task) => {
-          return client.fineTuning.jobs.cancel(params.id);
+          return client.fineTuning.jobs.cancel(params.id, { idempotencyKey: task.idempotencyKey });
         },
         {
           name: "Cancel Fine Tuning Job",

integrations/openai/src/images.ts

@@ -64,7 +64,7 @@ export class Images {
     return this.runTask(
       key,
       async (client, task) => {
-        return client.images.generate(params);
+        return client.images.generate(params, { idempotencyKey: task.idempotencyKey });
       },
       {
         name: "Create image",
@@ -121,15 +121,18 @@ export class Images {
           typeof params.image === "string" ? await fileFromUrl(params.image) : params.image;
         const mask = typeof params.mask === "string" ? await fileFromUrl(params.mask) : params.mask;
 
-        const response = await client.images.edit({
+        const response = await client.images.edit(
-          image: file,
+          {
-          prompt: params.prompt,
+            image: file,
-          mask: mask,
+            prompt: params.prompt,
-          n: params.n,
+            mask: mask,
-          size: params.size,
+            n: params.n,
-          response_format: params.response_format,
+            size: params.size,
-          user: params.user,
+            response_format: params.response_format,
-        });
+            user: params.user,
+          },
+          { idempotencyKey: task.idempotencyKey }
+        );
 
         return response;
       },
@@ -182,13 +185,16 @@ export class Images {
         const file =
           typeof params.image === "string" ? await fileFromUrl(params.image) : params.image;
 
-        const response = await client.images.createVariation({
+        const response = await client.images.createVariation(
-          image: file,
+          {
-          n: params.n,
+            image: file,
-          size: params.size,
+            n: params.n,
-          response_format: params.response_format,
+            size: params.size,
-          user: params.user,
+            response_format: params.response_format,
-        });
+            user: params.user,
+          },
+          { idempotencyKey: task.idempotencyKey }
+        );
 
         return response;
       },

integrations/openai/src/index.ts

@@ -55,6 +55,7 @@ export class OpenAI implements TriggerIntegration {
     this.native = new OpenAIApi({
       apiKey: options.apiKey,
       organization: options.organization,
+      baseURL: options.baseURL,
     });
   }
 
@@ -77,6 +78,7 @@ export class OpenAI implements TriggerIntegration {
     openai._client = new OpenAIApi({
       apiKey,
       organization: this._options.organization,
+      baseURL: this._options.baseURL ?? "https://api.openai.com/v1",
     });
     return openai;
   }
@@ -104,7 +106,7 @@ export class OpenAI implements TriggerIntegration {
         return callback(this._client, task, io);
       },
       {
-        icon: "openai",
+        icon: this._options.icon ?? "openai",
         retry: retry.standardBackoff,
         ...(options ?? {}),
         connectionKey: this._connectionKey,

integrations/openai/src/types.ts

@@ -2,6 +2,8 @@ export type OpenAIIntegrationOptions = {
   id: string;
   apiKey?: string;
   organization?: string;
+  baseURL?: string;
+  icon?: string;
 };
 
 export type OpenAIIntegrationAuth = Omit<OpenAIIntegrationOptions, "id">;

references/job-catalog/src/openai.ts

@@ -54,11 +54,6 @@ client.defineJob({
       ],
     });
 
-    await io.openai.backgroundCreateCompletion("background-completion", {
-      model: "text-davinci-003",
-      prompt: "Create a good programming joke about Tasks",
-    });
-
     await io.openai.createCompletion("completion", {
       model: "text-davinci-003",
       prompt: "Create a good programming joke about Tasks",
@@ -77,4 +72,44 @@ client.defineJob({
   },
 });
 
+const perplexity = new OpenAI({
+  id: "perplexity",
+  apiKey: process.env["PERPLEXITY_API_KEY"]!,
+  baseURL: "https://api.perplexity.ai",
+  icon: "brand-open-source",
+});
+
+client.defineJob({
+  id: "perplexity-tasks",
+  name: "Perplexity Tasks",
+  version: "0.0.1",
+  trigger: eventTrigger({
+    name: "perplexity.tasks",
+  }),
+  integrations: {
+    perplexity,
+  },
+  run: async (payload, io, ctx) => {
+    await io.perplexity.createChatCompletion("chat-completion", {
+      model: "mistral-7b-instruct",
+      messages: [
+        {
+          role: "user",
+          content: "Create a good programming joke about background jobs",
+        },
+      ],
+    });
+
+    await io.perplexity.backgroundCreateChatCompletion("background-chat-completion", {
+      model: "mistral-7b-instruct",
+      messages: [
+        {
+          role: "user",
+          content: "If you were a programming language, what would you be and why?",
+        },
+      ],
+    });
+  },
+});
+
 createExpressServer(client);