core[patch]: Fix formatting mustache image templates (#5666)

* core[patch]: Fix formatting mustache image templates * unfocus test * fix tests
langchain-ai · Jun 5, 2024 · e9a08aa · e9a08aa
1 parent e834086
commit e9a08aa
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 6 deletions.
diff --git a/langchain-core/src/prompts/chat.ts b/langchain-core/src/prompts/chat.ts
@@ -32,7 +32,12 @@ import {
   ExtractedFStringParams,
 } from "./prompt.js";
 import { ImagePromptTemplate } from "./image.js";
-import { TemplateFormat, parseFString } from "./template.js";
+import {
+  ParsedTemplateNode,
+  TemplateFormat,
+  parseFString,
+  parseMustache,
+} from "./template.js";
 
 /**
  * Abstract class that serves as a base for creating message prompt
@@ -495,13 +500,19 @@ class _StringImageMessagePromptTemplate<
         } else if (typeof item.text === "string") {
           text = item.text ?? "";
         }
-        prompt.push(PromptTemplate.fromTemplate(text));
+        prompt.push(PromptTemplate.fromTemplate(text, additionalOptions));
       } else if (typeof item === "object" && "image_url" in item) {
         let imgTemplate = item.image_url ?? "";
         let imgTemplateObject: ImagePromptTemplate<InputValues>;
         let inputVariables: string[] = [];
         if (typeof imgTemplate === "string") {
-          const parsedTemplate = parseFString(imgTemplate);
+          let parsedTemplate: ParsedTemplateNode[];
+          if (additionalOptions?.templateFormat === "mustache") {
+            parsedTemplate = parseMustache(imgTemplate);
+          } else {
+            parsedTemplate = parseFString(imgTemplate);
+          }
+
           const variables = parsedTemplate.flatMap((item) =>
             item.type === "variable" ? [item.name] : []
           );
@@ -524,7 +535,13 @@ class _StringImageMessagePromptTemplate<
           });
         } else if (typeof imgTemplate === "object") {
           if ("url" in imgTemplate) {
-            const parsedTemplate = parseFString(imgTemplate.url);
+            let parsedTemplate: ParsedTemplateNode[];
+            if (additionalOptions?.templateFormat === "mustache") {
+              parsedTemplate = parseMustache(imgTemplate.url);
+            } else {
+              parsedTemplate = parseFString(imgTemplate.url);
+            }
+
             inputVariables = parsedTemplate.flatMap((item) =>
               item.type === "variable" ? [item.name] : []
             );
@@ -913,7 +930,12 @@ export class ChatPromptTemplate<
           imageUrl = item.image_url.url;
         }
 
-        const promptTemplatePlaceholder = PromptTemplate.fromTemplate(imageUrl);
+        const promptTemplatePlaceholder = PromptTemplate.fromTemplate(
+          imageUrl,
+          {
+            templateFormat: this.templateFormat,
+          }
+        );
         const formattedUrl = await promptTemplatePlaceholder.format(
           inputValues
         );

diff --git a/langchain-core/src/prompts/tests/chat.mustache.test.ts b/langchain-core/src/prompts/tests/chat.mustache.test.ts
@@ -2,7 +2,7 @@ import { test, expect } from "@jest/globals";
 import { AIMessage } from "../../messages/ai.js";
 import { HumanMessage } from "../../messages/human.js";
 import { SystemMessage } from "../../messages/system.js";
-import { ChatPromptTemplate } from "../chat.js";
+import { ChatPromptTemplate, HumanMessagePromptTemplate } from "../chat.js";
 
 test("Test creating a chat prompt template from role string messages", async () => {
   const template = ChatPromptTemplate.fromMessages(
@@ -67,3 +67,53 @@ test("Ignores f-string inputs input variables with repeats.", async () => {
     new HumanMessage("This {bar} is a {foo} test {foo}."),
   ]);
 });
+
+test("Mustache template with image and chat prompts inside one template (fromMessages)", async () => {
+  const template = ChatPromptTemplate.fromMessages(
+    [
+      [
+        "human",
+        [
+          {
+            type: "image_url",
+            image_url: "{{image_url}}",
+          },
+          {
+            type: "text",
+            text: "{{other_var}}",
+          },
+        ],
+      ],
+      ["human", "hello {{name}}"],
+    ],
+    {
+      templateFormat: "mustache",
+    }
+  );
+
+  expect(template.inputVariables.sort()).toEqual([
+    "image_url",
+    "name",
+    "other_var",
+  ]);
+});
+
+test("Mustache image template with nested URL and chat prompts HumanMessagePromptTemplate.fromTemplate", async () => {
+  const template = HumanMessagePromptTemplate.fromTemplate(
+    [
+      {
+        text: "{{name}}",
+      },
+      {
+        image_url: {
+          url: "{{image_url}}",
+        },
+      },
+    ],
+    {
+      templateFormat: "mustache",
+    }
+  );
+
+  expect(template.inputVariables.sort()).toEqual(["image_url", "name"]);
+});