aimakerspace/openai_utils/prompts.py

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, format=True, **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
        """
        if format:
            return {"role": self.role, "content": self.format_prompt(**kwargs)}
        
        return {"role": self.role, "content": self.prompt}


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())