update checklist node definition

pufflyai · Nov 23, 2023 · 5a9950b · 5a9950b
1 parent 1f0c48b
commit 5a9950b
Show file tree

Hide file tree

Showing 5 changed files with 288 additions and 61 deletions.
diff --git a/packages/@pufflig/ps-nodes-config/src/adapters/document_check/document_check.ts b/packages/@pufflig/ps-nodes-config/src/adapters/document_check/document_check.ts
@@ -4,7 +4,7 @@ import { NodeConfig } from "@pufflig/ps-types";
 export const documentCheckNodeType = "modifier/document_check" as const;
 
 export const documentCheck: NodeConfig = {
-  name: "Document Checklist",
+  name: "Checklist",
   description: "Run a checklist on a document.",
   tags: ["adapter", "document", "text"],
   status: "stable",
@@ -23,9 +23,9 @@ export const documentCheck: NodeConfig = {
   },
   outputs: [
     {
-      id: "result",
-      name: "Result",
-      description: "A list, checklist or other information about the document",
+      id: "checklist",
+      name: "Checklist",
+      description: "A checklist of items to run on the document",
       type: "text",
       defaultValue: "",
     },
@@ -43,25 +43,44 @@ export const documentCheck: NodeConfig = {
       },
     },
     {
-      id: "prompt",
-      name: "Prompt",
-      description: "Prompt to check the document with",
+      id: "instructions",
+      name: "Instructions",
+      description: "Instructions for the AI",
       type: "text",
-      defaultValue: `Extract information in the document below and insert them in the csv table, don't overwrite existing values and keep things empty if you cannot find information in the document:\n\nTABLE EXAMPLE:\ncharacters, age\nmickey mouse, 10\ndonald duck, -\n\nTABLE:\n{{table}}\n\nDOCUMENT:\n{{document}}\n\nTABLE:\n`,
-    },
-    {
-      id: "table",
-      name: "Table",
-      description: "The list, table or checklist to parse the document with.",
-      type: "text",
-      defaultValue: "",
+      defaultValue: `Run the checklist below on the document.`,
     },
     {
       id: "document",
       name: "Document",
-      description: "Document to be processed",
+      description: "Document to be checked",
       type: "text",
       defaultValue: "",
     },
+    {
+      id: "checklist",
+      name: "Checklist",
+      description: "The checklist to run on the document",
+      type: "object",
+      editableSchema: true,
+      defaultValue: [],
+    },
+    {
+      id: "format",
+      name: "Format",
+      description: "The format in which to return the cheklist results",
+      type: "selection",
+      defaultValue: "markdown",
+      options: [
+        { id: "csv", name: "CSV" },
+        { id: "markdown", name: "Markdown" },
+      ],
+    },
+    {
+      id: "fields",
+      name: "Fields",
+      description: "Custom fields to include in the output for each checklist item",
+      type: "list",
+      defaultValue: ["ok"],
+    },
   ],
 };
diff --git a/...s/@pufflig/ps-nodes/src/adapters/document_check/__snapshots__/document_check.test.ts.snap b/...s/@pufflig/ps-nodes/src/adapters/document_check/__snapshots__/document_check.test.ts.snap
@@ -4,7 +4,7 @@ exports[`documentCheck should extract variables correctly 1`] = `
 [
   {
     "defaultValue": {
-      "modelId": "test_model",
+      "modelId": "gpt-3.5-turbo-instruct",
       "parameters": {},
     },
     "definition": {
@@ -213,32 +213,52 @@ exports[`documentCheck should extract variables correctly 1`] = `
     "type": "model",
   },
   {
-    "defaultValue": "Hello, {{myVariable}}!",
-    "description": "Prompt to check the document with",
-    "id": "prompt",
-    "name": "Prompt",
+    "defaultValue": "Run the checklist below on the document.",
+    "description": "Instructions for the AI",
+    "id": "instructions",
+    "name": "Instructions",
     "type": "text",
   },
   {
-    "defaultValue": "test_table",
-    "description": "The list, table or checklist to parse the document with.",
-    "id": "table",
-    "name": "Table",
-    "type": "text",
-  },
-  {
-    "defaultValue": "This is a test document.",
-    "description": "Document to be processed",
+    "defaultValue": "",
+    "description": "Document to be checked",
     "id": "document",
     "name": "Document",
     "type": "text",
   },
   {
-    "defaultValue": "myValue",
-    "description": "",
-    "id": "myVariable",
-    "name": "myVariable",
-    "type": "text",
+    "defaultValue": [],
+    "description": "The checklist to run on the document",
+    "editableSchema": true,
+    "id": "checklist",
+    "name": "Checklist",
+    "type": "object",
+  },
+  {
+    "defaultValue": "markdown",
+    "description": "The format in which to return the cheklist results",
+    "id": "format",
+    "name": "Format",
+    "options": [
+      {
+        "id": "csv",
+        "name": "CSV",
+      },
+      {
+        "id": "markdown",
+        "name": "Markdown",
+      },
+    ],
+    "type": "selection",
+  },
+  {
+    "defaultValue": [
+      "ok",
+    ],
+    "description": "Custom fields to include in the output for each checklist item",
+    "id": "fields",
+    "name": "Fields",
+    "type": "list",
   },
 ]
 `;
diff --git a/packages/@pufflig/ps-nodes/src/adapters/document_check/document_check.test.ts b/packages/@pufflig/ps-nodes/src/adapters/document_check/document_check.test.ts
@@ -1,5 +1,5 @@
-import { execute, getInputDefinition, LLMCompletionInput } from "./document_check";
 import axios from "axios";
+import { execute, getInputDefinition, LLMCompletionInput } from "./document_check";
 
 jest.mock("axios");
 
@@ -8,15 +8,58 @@ describe("documentCheck", () => {
     jest.resetAllMocks();
   });
 
-  it("should return the completion string", async () => {
+  it("should return the resulting checklist", async () => {
     const input: LLMCompletionInput = {
-      prompt: "Hello, world!",
+      instructions: "Hello, world!",
       model: {
         modelId: "test_model",
         parameters: {},
       },
       document: "This is a test document.",
-      table: "test_table",
+      checklist: [
+        { id: "is_greeting", defaultValue: "test_table", type: "text", name: "is_greeting", description: "" },
+      ],
+      fields: ["ok"],
+      format: "csv",
+    };
+
+    const expectedOutput = { result: "mock checklist" };
+    const mockedAxiosResponse = { data: expectedOutput };
+    (axios.post as jest.MockedFunction<typeof axios.post>).mockResolvedValueOnce(mockedAxiosResponse);
+
+    const output = await execute(input);
+
+    expect(output).toEqual({ checklist: expectedOutput.result });
+    expect(axios.post).toHaveBeenCalledTimes(1);
+  });
+
+  it("should parse input variables", async () => {
+    const input: LLMCompletionInput = {
+      instructions: "Hello, {{world}}! Run a checklist on the following document: {{document}}",
+      model: {
+        modelId: "test_model",
+        parameters: {},
+      },
+      document: "This is a test document.",
+      checklist: [
+        {
+          id: "is_greeting",
+          defaultValue: "is the text a greeting?",
+          type: "text",
+          name: "is_greeting",
+          description: "",
+        },
+        {
+          id: "is_formal",
+          defaultValue: "is the greeting formal?",
+          type: "text",
+          name: "is_formal",
+          description: "",
+        },
+      ],
+      fields: ["ok"],
+      format: "csv",
+      world: "test",
     };
 
     const expectedOutput = { result: "This is a test completion." };
@@ -25,20 +68,70 @@ describe("documentCheck", () => {
 
     const output = await execute(input);
 
-    expect(output).toEqual(expectedOutput);
+    expect(output).toEqual({ checklist: expectedOutput.result });
     expect(axios.post).toHaveBeenCalledTimes(1);
+    expect(axios.post).toHaveBeenCalledWith(
+      expect.any(String),
+      {
+        document: "This is a test document.",
+        format: `check,ok
+is_greeting,
+is_formal,`,
+        modelId: "test_model",
+        options: {
+          cache: true,
+          track: true,
+        },
+        parameters: {},
+        prompt: `Hello, test! Run a checklist on the following document: {{document}}
+
+CHECKLIST DESCRIPTION:
+check,description
+is_greeting, is the text a greeting?
+is_formal, is the greeting formal?
+
+CHECKLIST FORMAT:
+{{table}}
+
+CHECKLIST IN CSV FORMAT:
+`,
+      },
+      {
+        headers: {
+          Authorization: "Bearer undefined",
+          "Content-Type": "application/json",
+        },
+      }
+    );
   });
 
   it("should parse input variables", async () => {
     const input: LLMCompletionInput = {
-      prompt: "Hello, {{myVariable}}!",
+      instructions: "Hello, {{world}}! Run a checklist on the following document: {{document}}",
       model: {
         modelId: "test_model",
         parameters: {},
       },
       document: "This is a test document.",
-      table: "test_table",
-      myVariable: "myValue",
+      checklist: [
+        {
+          id: "is_greeting",
+          defaultValue: "is the text a greeting?",
+          type: "text",
+          name: "is_greeting",
+          description: "",
+        },
+        {
+          id: "is_formal",
+          defaultValue: "is the greeting formal?",
+          type: "text",
+          name: "is_formal",
+          description: "",
+        },
+      ],
+      fields: ["ok"],
+      format: "markdown",
+      world: "test",
     };
 
     const expectedOutput = { result: "This is a test completion." };
@@ -47,20 +140,33 @@ describe("documentCheck", () => {
 
     const output = await execute(input);
 
-    expect(output).toEqual(expectedOutput);
+    expect(output).toEqual({ checklist: expectedOutput.result });
     expect(axios.post).toHaveBeenCalledTimes(1);
     expect(axios.post).toHaveBeenCalledWith(
       expect.any(String),
       {
         document: "This is a test document.",
-        format: "test_table",
+        format: `|check|ok|
+|is_greeting||
+|is_formal||`,
         modelId: "test_model",
         options: {
           cache: true,
           track: true,
         },
         parameters: {},
-        prompt: "Hello, myValue!",
+        prompt: `Hello, test! Run a checklist on the following document: {{document}}
+
+CHECKLIST DESCRIPTION:
+|check|description|
+|is_greeting|is the text a greeting?|
+|is_formal|is the greeting formal?|
+
+CHECKLIST FORMAT:
+{{table}}
+
+CHECKLIST IN MARKDOWN FORMAT:
+`,
       },
       {
         headers: {
@@ -73,28 +179,49 @@ describe("documentCheck", () => {
 
   it("should extract variables correctly", async () => {
     const output = getInputDefinition({
-      prompt: "Hello, {{myVariable}}!",
+      instructions: "Hello, {{world}}!",
       model: {
         modelId: "test_model",
         parameters: {},
       },
       document: "This is a test document.",
-      table: "test_table",
-      myVariable: "myValue",
+      checklist: [
+        {
+          id: "is_greeting",
+          defaultValue: "is the text a greeting?",
+          type: "text",
+          name: "is_greeting",
+          description: "",
+        },
+        {
+          id: "is_formal",
+          defaultValue: "is the greeting formal?",
+          type: "text",
+          name: "is_formal",
+          description: "",
+        },
+      ],
+      fields: ["ok"],
+      format: "csv",
+      world: "test",
     });
 
     expect(output).toMatchSnapshot();
   });
 
   it("should throw an error if the API call fails", async () => {
     const input: LLMCompletionInput = {
-      prompt: "Hello, world!",
+      instructions: "Hello, world!",
       model: {
         modelId: "test_model",
         parameters: {},
       },
       document: "This is a test document.",
-      table: "test_table",
+      checklist: [
+        { id: "is_greeting", defaultValue: "test_table", type: "text", name: "is_greeting", description: "" },
+      ],
+      fields: ["ok"],
+      format: "csv",
     };
 
     const expectedError = new Error("API call failed.");