feat: optimize template parse (#460)

api/core/prompt/prompt_builder.py

@@ -3,13 +3,13 @@ import re
 from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, AIMessagePromptTemplate
 from langchain.schema import BaseMessage
 
-from core.prompt.prompt_template import OutLinePromptTemplate
+from core.prompt.prompt_template import JinjaPromptTemplate
 
 
 class PromptBuilder:
     @classmethod
     def to_system_message(cls, prompt_content: str, inputs: dict) -> BaseMessage:
-        prompt_template = OutLinePromptTemplate.from_template(prompt_content)
+        prompt_template = JinjaPromptTemplate.from_template(prompt_content)
         system_prompt_template = SystemMessagePromptTemplate(prompt=prompt_template)
         prompt_inputs = {k: inputs[k] for k in system_prompt_template.input_variables if k in inputs}
         system_message = system_prompt_template.format(**prompt_inputs)
@@ -17,7 +17,7 @@ class PromptBuilder:
 
     @classmethod
     def to_ai_message(cls, prompt_content: str, inputs: dict) -> BaseMessage:
-        prompt_template = OutLinePromptTemplate.from_template(prompt_content)
+        prompt_template = JinjaPromptTemplate.from_template(prompt_content)
         ai_prompt_template = AIMessagePromptTemplate(prompt=prompt_template)
         prompt_inputs = {k: inputs[k] for k in ai_prompt_template.input_variables if k in inputs}
         ai_message = ai_prompt_template.format(**prompt_inputs)
@@ -25,13 +25,14 @@ class PromptBuilder:
 
     @classmethod
     def to_human_message(cls, prompt_content: str, inputs: dict) -> BaseMessage:
-        prompt_template = OutLinePromptTemplate.from_template(prompt_content)
+        prompt_template = JinjaPromptTemplate.from_template(prompt_content)
         human_prompt_template = HumanMessagePromptTemplate(prompt=prompt_template)
         human_message = human_prompt_template.format(**inputs)
         return human_message
 
     @classmethod
     def process_template(cls, template: str):
-        processed_template = re.sub(r'\{([a-zA-Z_]\w+?)\}', r'\1', template)
-        processed_template = re.sub(r'\{\{([a-zA-Z_]\w+?)\}\}', r'{\1}', processed_template)
+        processed_template = re.sub(r'\{{2}(.+)\}{2}', r'{\1}', template)
+        # processed_template = re.sub(r'\{([a-zA-Z_]\w+?)\}', r'\1', template)
+        # processed_template = re.sub(r'\{\{([a-zA-Z_]\w+?)\}\}', r'{\1}', processed_template)
         return processed_template

api/core/prompt/prompt_template.py

@@ -1,10 +1,33 @@
 import re
 from typing import Any
 
+from jinja2 import Environment, meta
 from langchain import PromptTemplate
 from langchain.formatting import StrictFormatter
 
 
+class JinjaPromptTemplate(PromptTemplate):
+    template_format: str = "jinja2"
+    """The format of the prompt template. Options are: 'f-string', 'jinja2'."""
+
+    @classmethod
+    def from_template(cls, template: str, **kwargs: Any) -> PromptTemplate:
+        """Load a prompt template from a template."""
+        env = Environment()
+        ast = env.parse(template)
+        input_variables = meta.find_undeclared_variables(ast)
+
+        if "partial_variables" in kwargs:
+            partial_variables = kwargs["partial_variables"]
+            input_variables = {
+                var for var in input_variables if var not in partial_variables
+            }
+
+        return cls(
+            input_variables=list(sorted(input_variables)), template=template, **kwargs
+        )
+
+
 class OutLinePromptTemplate(PromptTemplate):
     @classmethod
     def from_template(cls, template: str, **kwargs: Any) -> PromptTemplate:
@@ -16,6 +39,24 @@ class OutLinePromptTemplate(PromptTemplate):
             input_variables=list(sorted(input_variables)), template=template, **kwargs
         )
 
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt with the inputs.
+
+        Args:
+            kwargs: Any arguments to be passed to the prompt template.
+
+        Returns:
+            A formatted string.
+
+        Example:
+
+        .. code-block:: python
+
+            prompt.format(variable1="foo")
+        """
+        kwargs = self._merge_partial_and_user_variables(**kwargs)
+        return OneLineFormatter().format(self.template, **kwargs)
+
 
 class OneLineFormatter(StrictFormatter):
     def parse(self, format_string):

api/core/prompt/prompts.py

@@ -1,5 +1,5 @@
 CONVERSATION_TITLE_PROMPT = (
-    "Human:{query}\n-----\n"
+    "Human:{{query}}\n-----\n"
     "Help me summarize the intent of what the human said and provide a title, the title should not exceed 20 words.\n"
     "If the human said is conducted in Chinese, you should return a Chinese title.\n" 
     "If the human said is conducted in English, you should return an English title.\n"
@@ -19,7 +19,7 @@ CONVERSATION_SUMMARY_PROMPT = (
 INTRODUCTION_GENERATE_PROMPT = (
     "I am designing a product for users to interact with an AI through dialogue. "
     "The Prompt given to the AI before the conversation is:\n\n"
-    "```\n{prompt}\n```\n\n"
+    "```\n{{prompt}}\n```\n\n"
     "Please generate a brief introduction of no more than 50 words that greets the user, based on this Prompt. "
     "Do not reveal the developer's motivation or deep logic behind the Prompt, "
     "but focus on building a relationship with the user:\n"
@@ -27,13 +27,13 @@ INTRODUCTION_GENERATE_PROMPT = (
 
 MORE_LIKE_THIS_GENERATE_PROMPT = (
     "-----\n"
-    "{original_completion}\n"
+    "{{original_completion}}\n"
     "-----\n\n"
     "Please use the above content as a sample for generating the result, "
     "and include key information points related to the original sample in the result. "
     "Try to rephrase this information in different ways and predict according to the rules below.\n\n"
     "-----\n"
-    "{prompt}\n"
+    "{{prompt}}\n"
 )
 
 SUGGESTED_QUESTIONS_AFTER_ANSWER_INSTRUCTION_PROMPT = (