File size: 2,469 Bytes
234eac0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
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())
|