dify/api/core/prompt/prompt_template.py

40 lines
1.5 KiB
Python
Raw Normal View History

2023-05-15 08:51:32 +08:00
import re
REGEX = re.compile(r"\{\{([a-zA-Z_][a-zA-Z0-9_]{0,29}|#histories#|#query#|#context#)\}\}")
2023-05-15 08:51:32 +08:00
class PromptTemplateParser:
"""
Rules:
2023-06-27 15:30:38 +08:00
1. Template variables must be enclosed in `{{}}`.
2. The template variable Key can only be: letters + numbers + underscore, with a maximum length of 16 characters,
and can only start with letters and underscores.
3. The template variable Key cannot contain new lines or spaces, and must comply with rule 2.
4. In addition to the above, 3 types of special template variable Keys are accepted:
`{{#histories#}}` `{{#query#}}` `{{#context#}}`. No other `{{##}}` template variables are allowed.
"""
2023-06-27 15:30:38 +08:00
def __init__(self, template: str):
self.template = template
self.variable_keys = self.extract()
2023-06-27 15:30:38 +08:00
def extract(self) -> list:
# Regular expression to match the template rules
return re.findall(REGEX, self.template)
2023-06-27 15:30:38 +08:00
def format(self, inputs: dict, remove_template_variables: bool = True) -> str:
def replacer(match):
key = match.group(1)
value = inputs.get(key, match.group(0)) # return original matched string if key not found
2023-06-27 15:30:38 +08:00
if remove_template_variables:
return PromptTemplateParser.remove_template_variables(value)
return value
2023-06-27 15:30:38 +08:00
return re.sub(REGEX, replacer, self.template)
2023-06-27 15:30:38 +08:00
@classmethod
def remove_template_variables(cls, text: str):
return re.sub(REGEX, r'{\1}', text)