itsprasun/pythonic-rag-fastapi-react

1import re2 3 4class BasePrompt:5    def __init__(self, prompt):6        """7        Initializes the BasePrompt object with a prompt template.8 9        :param prompt: A string that can contain placeholders within curly braces10        """11        self.prompt = prompt12        self._pattern = re.compile(r"\{([^}]+)\}")13 14    def format_prompt(self, **kwargs):15        """16        Formats the prompt string using the keyword arguments provided.17 18        :param kwargs: The values to substitute into the prompt string19        :return: The formatted prompt string20        """21        matches = self._pattern.findall(self.prompt)22        return self.prompt.format(**{match: kwargs.get(match, "") for match in matches})23 24    def get_input_variables(self):25        """26        Gets the list of input variable names from the prompt string.27 28        :return: List of input variable names29        """30        return self._pattern.findall(self.prompt)31 32 33class RolePrompt(BasePrompt):34    def __init__(self, prompt, role: str):35        """36        Initializes the RolePrompt object with a prompt template and a role.37 38        :param prompt: A string that can contain placeholders within curly braces39        :param role: The role for the message ('system', 'user', or 'assistant')40        """41        super().__init__(prompt)42        self.role = role43 44    def create_message(self, format=True, **kwargs):45        """46        Creates a message dictionary with a role and a formatted message.47 48        :param kwargs: The values to substitute into the prompt string49        :return: Dictionary containing the role and the formatted message50        """51        if format:52            return {"role": self.role, "content": self.format_prompt(**kwargs)}53        54        return {"role": self.role, "content": self.prompt}55 56 57class SystemRolePrompt(RolePrompt):58    def __init__(self, prompt: str):59        super().__init__(prompt, "system")60 61 62class UserRolePrompt(RolePrompt):63    def __init__(self, prompt: str):64        super().__init__(prompt, "user")65 66 67class AssistantRolePrompt(RolePrompt):68    def __init__(self, prompt: str):69        super().__init__(prompt, "assistant")70 71 72if __name__ == "__main__":73    prompt = BasePrompt("Hello {name}, you are {age} years old")74    print(prompt.format_prompt(name="John", age=30))75 76    prompt = SystemRolePrompt("Hello {name}, you are {age} years old")77    print(prompt.create_message(name="John", age=30))78    print(prompt.get_input_variables())79