OpenAI: createEmbedding

matt-aitken · matt-aitken · commit 52aacc1e1d80 · 2023-06-22T16:29:00.000+01:00
diff --git a/examples/nextjs-example/src/jobs/openai.ts b/examples/nextjs-example/src/jobs/openai.ts
@@ -64,6 +64,11 @@ new Job(client, {
       input: "Thsi is ridddled with erors",
       instruction: "Fix the spelling errors",
     });
+
+    await io.openai.createEmbedding("embedding", {
+      model: "text-embedding-ada-002",
+      input: "The food was delicious and the waiter...",
+    });
   },
 });
 
diff --git a/integrations/openai/src/index.ts b/integrations/openai/src/index.ts
@@ -7,6 +7,7 @@ import {
   createChatCompletion,
   createCompletion,
   createEdit,
+  createEmbedding,
   createFile,
   createFineTune,
   createFineTuneFile,
@@ -30,6 +31,7 @@ const tasks = {
   backgroundCreateChatCompletion,
   createEdit,
   createImage,
+  createEmbedding,
   createFile,
   listFiles,
   createFineTuneFile,
diff --git a/integrations/openai/src/tasks.ts b/integrations/openai/src/tasks.ts
@@ -3,6 +3,7 @@ import {
   CreateChatCompletionRequest,
   CreateCompletionRequest,
   CreateEditRequest,
+  CreateEmbeddingRequest,
   CreateFineTuneRequest,
   CreateImageRequest,
   OpenAIApi,
@@ -268,7 +269,7 @@ export const createImage: AuthenticatedTask<
   Prettify<CreateImageRequest>,
   CreateImageResponseData
 > = {
-  run: async (params, client) => {
+  run: async (params, client, task) => {
     return client.createImage(params).then((res) => res.data);
   },
   init: (params) => {
@@ -309,6 +310,40 @@ export const createImage: AuthenticatedTask<
   },
 };
 
+type CreateEmbeddingResponseData = Prettify<
+  Awaited<ReturnType<OpenAIClientType["createEmbedding"]>>["data"]
+>;
+
+export const createEmbedding: AuthenticatedTask<
+  OpenAIClientType,
+  Prettify<CreateEmbeddingRequest>,
+  CreateEmbeddingResponseData
+> = {
+  run: async (params, client) => {
+    return client.createEmbedding(params).then((res) => res.data);
+  },
+  init: (params) => {
+    return {
+      name: "Create embedding",
+      params,
+      icon: "openai",
+      properties: [
+        {
+          label: "Model",
+          text: params.model,
+        },
+        {
+          label: "Input",
+          text:
+            typeof params.input === "string"
+              ? truncate(params.input, 40)
+              : truncate(params.input.at(0) ?? "none", 40),
+        },
+      ],
+    };
+  },
+};
+
 type CreateFileResponseData = Awaited<
   ReturnType<OpenAIClientType["createFile"]>
 >["data"];
