prompt_template.py

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}
        ]