Create prompt_template.py

prompt_template.py

@@ -0,0 +1,86 @@
+import re
+import yaml
+from dataclasses import dataclass, field
+from typing import Dict, List, Set, Optional
+
+from openai import OpenAI
+
+@dataclass
+class PromptTemplate:
+    """
+    A template class for managing and validating LLM prompts.
+
+    This class handles:
+    - Storing system and user prompts
+    - Validating required template variables
+    - Formatting prompts with provided variables
+
+    Attributes:
+        system_prompt (str): The system-level instructions for the LLM
+        user_template (str): Template string with variables in {variable} format
+    """
+    system_prompt: str
+    user_template: str
+
+    def __post_init__(self):
+        """Initialize the set of required variables from the template."""
+        self.required_variables: Set[str] = self._get_required_variables()
+
+    def _get_required_variables(self) -> set:
+        """
+        Extract required variables from the template using regex.
+
+        Returns:
+            set: Set of variable names found in the template
+
+        Example:
+            Template "Write about {topic} in {style}" returns {'topic', 'style'}
+        """
+        return set(re.findall(r'\{(\w+)\}', self.user_template))
+
+    def _validate_variables(self, provided_vars: Dict):
+        """
+        Ensure all required template variables are provided.
+
+        Args:
+            provided_vars: Dictionary of variable names and values
+
+        Raises:
+            ValueError: If any required variables are missing
+        """
+        provided_keys = set(provided_vars.keys())
+        missing_vars = self.required_variables - provided_keys
+        if missing_vars:
+            error_msg = (
+                f"\nPrompt Template Error:\n"
+                f"Missing required variables: {', '.join(missing_vars)}\n"
+                f"Template requires: {', '.join(self.required_variables)}\n"
+                f"You provided: {', '.join(provided_keys)}\n"
+                f"Template string: '{self.user_template}'"
+            )
+            raise ValueError(error_msg)
+
+    def format(self, **kwargs) -> List[Dict[str, str]]:
+        """
+        Format the prompt template with provided variables.
+
+        Args:
+            **kwargs: Key-value pairs for template variables
+
+        Returns:
+            List[Dict[str, str]]: Formatted messages ready for LLM API
+
+        Example:
+            template.format(topic="AI", style="academic")
+        """
+        self._validate_variables(kwargs)
+
+        try:
+            formatted_user_message = self.user_template.format(**kwargs)
+        except Exception as e:
+            raise ValueError(f"Error formatting template: {str(e)}")
+
+        return [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": formatted_user_message}
+        ]