A first step towards composable prompts: composable chat prompt templ…

…ates (langchain-ai#270) * A first step towards composable prompts: chat prompts * Implement partial * Implement support for partila
philrwebb · Apr 26, 2023 · 195767d · 195767d
1 parent 00bb72a
commit 195767d
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 4 deletions.
diff --git a/langchain/src/prompts/chat.ts b/langchain/src/prompts/chat.ts
@@ -266,7 +266,7 @@ export class ChatPromptTemplate
     };
   }
 
-  async partial(values: PartialValues): Promise<BasePromptTemplate> {
+  async partial(values: PartialValues): Promise<ChatPromptTemplate> {
     // This is implemented in a way it doesn't require making
     // BaseMessagePromptTemplate aware of .partial()
     const promptDict: ChatPromptTemplateInput = { ...this };
@@ -281,17 +281,40 @@ export class ChatPromptTemplate
   }
 
   static fromPromptMessages(
-    promptMessages: BaseMessagePromptTemplate[]
+    promptMessages: (BaseMessagePromptTemplate | ChatPromptTemplate)[]
   ): ChatPromptTemplate {
+    const flattenedMessages = promptMessages.reduce(
+      (acc, promptMessage) =>
+        acc.concat(
+          // eslint-disable-next-line no-instanceof/no-instanceof
+          promptMessage instanceof ChatPromptTemplate
+            ? promptMessage.promptMessages
+            : [promptMessage]
+        ),
+      [] as BaseMessagePromptTemplate[]
+    );
+    const flattenedPartialVariables = promptMessages.reduce(
+      (acc, promptMessage) =>
+        // eslint-disable-next-line no-instanceof/no-instanceof
+        promptMessage instanceof ChatPromptTemplate
+          ? Object.assign(acc, promptMessage.partialVariables)
+          : acc,
+      Object.create(null) as PartialValues
+    );
+
     const inputVariables = new Set<string>();
-    for (const promptMessage of promptMessages) {
+    for (const promptMessage of flattenedMessages) {
       for (const inputVariable of promptMessage.inputVariables) {
+        if (inputVariable in flattenedPartialVariables) {
+          continue;
+        }
         inputVariables.add(inputVariable);
       }
     }
     return new ChatPromptTemplate({
       inputVariables: [...inputVariables],
-      promptMessages,
+      promptMessages: flattenedMessages,
+      partialVariables: flattenedPartialVariables,
     });
   }
 }
diff --git a/langchain/src/prompts/tests/chat.test.ts b/langchain/src/prompts/tests/chat.test.ts
@@ -136,6 +136,67 @@ test("Test fromPromptMessages", async () => {
   ]);
 });
 
+test("Test fromPromptMessages is composable", async () => {
+  const systemPrompt = new PromptTemplate({
+    template: "Here's some context: {context}",
+    inputVariables: ["context"],
+  });
+  const userPrompt = new PromptTemplate({
+    template: "Hello {foo}, I'm {bar}",
+    inputVariables: ["foo", "bar"],
+  });
+  const chatPromptInner = ChatPromptTemplate.fromPromptMessages([
+    new SystemMessagePromptTemplate(systemPrompt),
+    new HumanMessagePromptTemplate(userPrompt),
+  ]);
+  const chatPrompt = ChatPromptTemplate.fromPromptMessages([
+    chatPromptInner,
+    AIMessagePromptTemplate.fromTemplate("I'm an AI. I'm {foo}. I'm {bar}."),
+  ]);
+  expect(chatPrompt.inputVariables).toEqual(["context", "foo", "bar"]);
+  const messages = await chatPrompt.formatPromptValue({
+    context: "This is a context",
+    foo: "Foo",
+    bar: "Bar",
+  });
+  expect(messages.toChatMessages()).toEqual([
+    new SystemChatMessage("Here's some context: This is a context"),
+    new HumanChatMessage("Hello Foo, I'm Bar"),
+    new AIChatMessage("I'm an AI. I'm Foo. I'm Bar."),
+  ]);
+});
+
+test("Test fromPromptMessages is composable with partial vars", async () => {
+  const systemPrompt = new PromptTemplate({
+    template: "Here's some context: {context}",
+    inputVariables: ["context"],
+  });
+  const userPrompt = new PromptTemplate({
+    template: "Hello {foo}, I'm {bar}",
+    inputVariables: ["foo", "bar"],
+  });
+  const chatPromptInner = ChatPromptTemplate.fromPromptMessages([
+    new SystemMessagePromptTemplate(systemPrompt),
+    new HumanMessagePromptTemplate(userPrompt),
+  ]);
+  const chatPrompt = ChatPromptTemplate.fromPromptMessages([
+    await chatPromptInner.partial({
+      context: "This is a context",
+      foo: "Foo",
+    }),
+    AIMessagePromptTemplate.fromTemplate("I'm an AI. I'm {foo}. I'm {bar}."),
+  ]);
+  expect(chatPrompt.inputVariables).toEqual(["bar"]);
+  const messages = await chatPrompt.formatPromptValue({
+    bar: "Bar",
+  });
+  expect(messages.toChatMessages()).toEqual([
+    new SystemChatMessage("Here's some context: This is a context"),
+    new HumanChatMessage("Hello Foo, I'm Bar"),
+    new AIChatMessage("I'm an AI. I'm Foo. I'm Bar."),
+  ]);
+});
+
 test("Test SimpleMessagePromptTemplate", async () => {
   const prompt = new MessagesPlaceholder("foo");
   const values = { foo: [new HumanChatMessage("Hello Foo, I'm Bar")] };