import re class BasePrompt: def __init__(self, prompt): """ Initializes the BasePrompt object with a prompt template. :param prompt: A string that can contain placeholders within curly braces """ self.prompt = prompt self._pattern = re.compile(r"\{([^}]+)\}") def format_prompt(self, **kwargs): """ Formats the prompt string using the keyword arguments provided. :param kwargs: The values to substitute into the prompt string :return: The formatted prompt string """ matches = self._pattern.findall(self.prompt) return self.prompt.format(**{match: kwargs.get(match, "") for match in matches}) def get_input_variables(self): """ Gets the list of input variable names from the prompt string. :return: List of input variable names """ return self._pattern.findall(self.prompt) class RolePrompt(BasePrompt): def __init__(self, prompt, role: str): """ Initializes the RolePrompt object with a prompt template and a role. :param prompt: A string that can contain placeholders within curly braces :param role: The role for the message ('system', 'user', or 'assistant') """ super().__init__(prompt) self.role = role def create_message(self, **kwargs): """ Creates a message dictionary with a role and a formatted message. :param kwargs: The values to substitute into the prompt string :return: Dictionary containing the role and the formatted message """ return {"role": self.role, "content": self.format_prompt(**kwargs)} class SystemRolePrompt(RolePrompt): def __init__(self, prompt: str): super().__init__(prompt, "system") class UserRolePrompt(RolePrompt): def __init__(self, prompt: str): super().__init__(prompt, "user") class AssistantRolePrompt(RolePrompt): def __init__(self, prompt: str): super().__init__(prompt, "assistant") if __name__ == "__main__": prompt = BasePrompt("Hello {name}, you are {age} years old") print(prompt.format_prompt(name="John", age=30)) prompt = SystemRolePrompt("Hello {name}, you are {age} years old") print(prompt.create_message(name="John", age=30)) print(prompt.get_input_variables())