fix variable format in the document check node

pufflyai · Oct 31, 2023 · 6065461 · 6065461
1 parent 3787827
commit 6065461
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 4 deletions.
diff --git a/packages/@pufflig/ps-nodes-config/src/modifiers/document_check/document_check.ts b/packages/@pufflig/ps-nodes-config/src/modifiers/document_check/document_check.ts
@@ -47,7 +47,7 @@ export const documentCheck: NodeConfig = {
       name: "Prompt",
       description: "Prompt to check the document with",
       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`,
+      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",

diff --git a/packages/@pufflig/ps-nodes/src/data/prompt/prompt.test.ts b/packages/@pufflig/ps-nodes/src/data/prompt/prompt.test.ts
@@ -6,7 +6,19 @@ test("execute - no variables", async () => {
   });
   expect(variables).toMatchInlineSnapshot(`
     {
-      "prompt": "summarize ",
+      "prompt": "summarize {{longText}}",
+    }
+  `);
+});
+
+test("execute - empty variables", async () => {
+  const variables = await execute({
+    template: `summarize {{longText}}`,
+    longText: "",
+  });
+  expect(variables).toMatchInlineSnapshot(`
+    {
+      "prompt": "summarize {{longText}}",
     }
   `);
 });

diff --git a/packages/@pufflig/ps-nodes/src/data/prompt/prompt.ts b/packages/@pufflig/ps-nodes/src/data/prompt/prompt.ts
@@ -14,7 +14,18 @@ export interface TemplateTextOutput {
 
 export const execute = async (input: TemplateTextInput) => {
   const { template, ...variables } = input;
-  const renderedTemplate = Mustache.render(template, variables);
+
+  // keep template values for undefined variables
+  const extractedVariables = extractVariables(template) || [];
+  const variablesWithDefaults = extractedVariables.reduce((acc, param) => {
+    return {
+      ...acc,
+      [param.id]:
+        variables[param.id] === undefined || variables[param.id] === "" ? `{{${param.id}}}` : variables[param.id],
+    };
+  }, {} as Record<string, string>);
+
+  const renderedTemplate = Mustache.render(template, variablesWithDefaults);
   return {
     prompt: renderedTemplate,
   };

diff --git a/packages/@pufflig/ps-nodes/src/modifiers/document_check/document_check.ts b/packages/@pufflig/ps-nodes/src/modifiers/document_check/document_check.ts
@@ -22,7 +22,8 @@ export const execute: Execute<LLMCompletionInput, LLMCompletionOutput> = async (
   const { modelId, parameters } = model;
   const { globals } = options;
 
-  const renderedPrompt = Mustache.render(prompt, variables);
+  // render the prompt without overwriting the document and table variables
+  const renderedPrompt = Mustache.render(prompt, { ...variables, document: "{{document}}", table: "{{table}}" });
 
   const { result } = await refineCompletion({
     apiKey: getPromptStudioKey(globals || {}),