add fiter option

hiroki0525 · Nov 4, 2023 · 04a30fb · 04a30fb
1 parent 002e5d1
commit 04a30fb
Show file tree

Hide file tree

Showing 2 changed files with 165 additions and 98 deletions.
diff --git a/packages/core/src/__tests__/index.test.ts b/packages/core/src/__tests__/index.test.ts
@@ -1,5 +1,10 @@
 import { mkdir, rm, rmdir, writeFile } from "fs/promises";
-import generateDandoriTasks, { DandoriTask } from "../index";
+import generateDandoriTasks, {
+  ChatGPTFunctionCallModel,
+  DandoriTaskOptionalProperty,
+  DandoriTaskProperty,
+  DandoriTaskRequiredProperty,
+} from "../index";
 import {
   describe,
   beforeEach,
@@ -184,12 +189,50 @@ describe("generateDandoriTasks", () => {
         },
       },
     } as const;
-    const requiredProperties: readonly (keyof DandoriTask)[] = [
+    const requiredProperties: readonly DandoriTaskRequiredProperty[] = [
       "id",
       "name",
       "fromTaskIdList",
     ];
     const logPrefix = "Generating tasks";
+    const createOpenAiChatGptArguments = ({
+      source,
+      model = "gpt-3.5-turbo-0613",
+      filter = Object.keys(functionCallTaskProperties) as DandoriTaskProperty[],
+    }: {
+      source: string;
+      model?: ChatGPTFunctionCallModel;
+      filter?: DandoriTaskProperty[];
+    }) => ({
+      messages: [{ role: "user", content: source }],
+      model,
+      function_call: { name: functionCallName },
+      functions: [
+        {
+          name: functionCallName,
+          description:
+            "Get the tasks flow which will be used like Gantt chart.",
+          parameters: {
+            type: "object",
+            properties: {
+              tasks: {
+                type: "array",
+                items: {
+                  type: "object",
+                  required: requiredProperties,
+                  properties: Object.fromEntries(
+                    filter.map((prop) => [
+                      prop,
+                      functionCallTaskProperties[prop],
+                    ]),
+                  ),
+                },
+              },
+            },
+          },
+        },
+      ],
+    });
 
     beforeEach(() => {
       process.env[openApiKeyPropName] = apiKey;
@@ -199,7 +242,7 @@ describe("generateDandoriTasks", () => {
       expect(process.env[openApiKeyPropName]).toBe(apiKey);
     });
 
-    describe("with model argument", () => {
+    describe("with options which include model argument", () => {
       let result: Awaited<ReturnType<typeof generateDandoriTasks>>;
       const source = "with model argument";
       const model = "gpt-4-0613";
@@ -211,33 +254,46 @@ describe("generateDandoriTasks", () => {
       });
 
       it("called chat.completions.create with valid arguments", () => {
-        expect(openAI.chat.completions.create).toBeCalledWith({
-          messages: [{ role: "user", content: source }],
-          model,
-          function_call: { name: functionCallName },
-          functions: [
-            {
-              name: functionCallName,
-              description:
-                "Get the tasks flow which will be used like Gantt chart.",
-              parameters: {
-                type: "object",
-                properties: {
-                  tasks: {
-                    type: "array",
-                    items: {
-                      type: "object",
-                      required: requiredProperties,
-                      properties: functionCallTaskProperties,
-                    },
-                  },
-                },
-              },
-            },
-          ],
+        expect(openAI.chat.completions.create).toBeCalledWith(
+          createOpenAiChatGptArguments({ source, model }),
+        );
+      });
+
+      it("called logger.debug with valid arguments", () => {
+        expect(logger.debug).toBeCalledWith(openAiResArguments.tasks);
+      });
+
+      it("return tasks", () => {
+        expect(result).toStrictEqual(openAiResArguments.tasks);
+      });
+
+      it("called log with valid statement", () => {
+        expect(runPromisesSequentiallyMock.mock.calls[0][1]).toContain(
+          logPrefix,
+        );
+      });
+    });
+
+    describe("with options which include filter argument", () => {
+      let result: Awaited<ReturnType<typeof generateDandoriTasks>>;
+      const source = "with filter argument";
+      const filter: DandoriTaskOptionalProperty[] = ["deadline"];
+
+      beforeEach(async () => {
+        result = await generateDandoriTasks(source, {
+          filter,
         });
       });
 
+      it("called chat.completions.create with valid arguments", () => {
+        expect(openAI.chat.completions.create).toBeCalledWith(
+          createOpenAiChatGptArguments({
+            source,
+            filter: [...filter, ...requiredProperties],
+          }),
+        );
+      });
+
       it("called logger.debug with valid arguments", () => {
         expect(logger.debug).toBeCalledWith(openAiResArguments.tasks);
       });
@@ -253,7 +309,7 @@ describe("generateDandoriTasks", () => {
       });
     });
 
-    describe("without model argument", () => {
+    describe("without options", () => {
       let result: Awaited<ReturnType<typeof generateDandoriTasks>>;
       const source = "without model argument";
 
@@ -262,31 +318,9 @@ describe("generateDandoriTasks", () => {
       });
 
       it("called chat.completions.create with valid arguments", () => {
-        expect(openAI.chat.completions.create).toBeCalledWith({
-          messages: [{ role: "user", content: source }],
-          model: "gpt-3.5-turbo-0613",
-          function_call: { name: functionCallName },
-          functions: [
-            {
-              name: functionCallName,
-              description:
-                "Get the tasks flow which will be used like Gantt chart.",
-              parameters: {
-                type: "object",
-                properties: {
-                  tasks: {
-                    type: "array",
-                    items: {
-                      type: "object",
-                      required: requiredProperties,
-                      properties: functionCallTaskProperties,
-                    },
-                  },
-                },
-              },
-            },
-          ],
-        });
+        expect(openAI.chat.completions.create).toBeCalledWith(
+          createOpenAiChatGptArguments({ source }),
+        );
       });
 
       it("called logger.debug with valid arguments", () => {

diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
@@ -12,11 +12,6 @@ export type ChatGPTFunctionCallModel = "gpt-3.5-turbo-0613" | "gpt-4-0613";
 const defaultChatGPTFunctionCallModel: ChatGPTFunctionCallModel =
   "gpt-3.5-turbo-0613";
 
-export type GenerateDandoriTasksOptions = {
-  chatGPTModel?: ChatGPTFunctionCallModel;
-  envFilePath?: string;
-};
-
 export type DandoriTask = {
   id: string;
   name: string;
@@ -29,6 +24,24 @@ export type DandoriTask = {
   fromTaskIdList: string[];
 };
 
+export type DandoriTaskProperty = keyof DandoriTask;
+
+export type DandoriTaskRequiredProperty = Extract<
+  DandoriTaskProperty,
+  "id" | "name" | "fromTaskIdList"
+>;
+
+export type DandoriTaskOptionalProperty = Exclude<
+  DandoriTaskProperty,
+  DandoriTaskRequiredProperty
+>;
+
+export type GenerateDandoriTasksOptions = {
+  chatGPTModel?: ChatGPTFunctionCallModel;
+  envFilePath?: string;
+  filter?: DandoriTaskOptionalProperty[];
+};
+
 type FunctionCallValue = {
   type: "string" | "array" | "object";
   description?: string;
@@ -40,52 +53,62 @@ const excludePropertyPrompt =
   "If not provided, this property shouldn't be included.";
 const generateIdPrompt = "If not provided, return generated unique ID.";
 
-const functionCallTaskProperties: Record<keyof DandoriTask, FunctionCallValue> =
-  {
-    id: {
+const requiredFunctionCallTaskProperties: Record<
+  DandoriTaskRequiredProperty,
+  FunctionCallValue
+> = {
+  id: {
+    type: "string",
+    description: `The task ID. ${generateIdPrompt}`,
+  },
+  name: {
+    type: "string",
+    description: "The task name",
+  },
+  fromTaskIdList: {
+    type: "array",
+    description: "Task IDs to be executed before this task",
+    items: {
       type: "string",
-      description: `The task ID. ${generateIdPrompt}`,
     },
-    name: {
-      type: "string",
-      description: "The task name",
-    },
-    description: {
-      type: "string",
-      description: "The task description",
-    },
-    deadline: {
-      type: "string",
-      description: `The task deadline which is used by JavaScript Date constructor arguments. ${excludePropertyPrompt}`,
-    },
-    assignee: {
-      type: "object",
-      description: `The task assignee. ${excludePropertyPrompt}`,
-      properties: {
-        id: {
-          type: "string",
-          description: `The task assignee ID. ${generateIdPrompt}`,
-        },
-        name: {
-          type: "string",
-          description: "The task assignee name.",
-        },
+  },
+} as const;
+
+const optionalFunctionCallTaskProperties: Record<
+  DandoriTaskOptionalProperty,
+  FunctionCallValue
+> = {
+  description: {
+    type: "string",
+    description: "The task description",
+  },
+  deadline: {
+    type: "string",
+    description: `The task deadline which is used by JavaScript Date constructor arguments. ${excludePropertyPrompt}`,
+  },
+  assignee: {
+    type: "object",
+    description: `The task assignee. ${excludePropertyPrompt}`,
+    properties: {
+      id: {
+        type: "string",
+        description: `The task assignee ID. ${generateIdPrompt}`,
       },
-    },
-    fromTaskIdList: {
-      type: "array",
-      description: "Task IDs to be executed before this task",
-      items: {
+      name: {
         type: "string",
+        description: "The task assignee name.",
       },
     },
-  } as const;
+  },
+};
+
+const requiredProperties: readonly DandoriTaskRequiredProperty[] = Object.keys(
+  requiredFunctionCallTaskProperties,
+) as DandoriTaskRequiredProperty[];
 
-const requiredProperties: readonly (keyof DandoriTask)[] = [
-  "id",
-  "name",
-  "fromTaskIdList",
-];
+const optionalProperties: readonly DandoriTaskOptionalProperty[] = Object.keys(
+  optionalFunctionCallTaskProperties,
+) as DandoriTaskOptionalProperty[];
 
 const functionCallName = "get_tasks_flow";
 
@@ -105,6 +128,13 @@ export default async function generateDandoriTasks(
   const openai = new OpenAI();
   const model: ChatGPTFunctionCallModel =
     options?.chatGPTModel ?? defaultChatGPTFunctionCallModel;
+  const filterProperties = options?.filter ?? optionalProperties;
+  const filteredOptionalFunctionCallTaskProperties = Object.fromEntries(
+    filterProperties.map((filterProperty) => [
+      filterProperty,
+      optionalFunctionCallTaskProperties[filterProperty],
+    ]),
+  );
   const [completion] = await runPromisesSequentially(
     [
       () =>
@@ -125,7 +155,10 @@ export default async function generateDandoriTasks(
                     items: {
                       type: "object",
                       required: requiredProperties,
-                      properties: functionCallTaskProperties,
+                      properties: {
+                        ...requiredFunctionCallTaskProperties,
+                        ...filteredOptionalFunctionCallTaskProperties,
+                      },
                     },
                   },
                 },