prompt¶

class BasePromptTemplate(BaseModel):
    """
    Abstract base class for prompt templates.

    This class provides a common interface for messages from template strings with variable substitutions.    

    Attributes
    ----------
    template_str : str
        The template string containing placeholders for variable substitution.
        The specific placeholder syntax depends on the concrete implementation
        (e.g., f-string, Jinja2, etc.).

    Methods
    -------
    format_message(role, **kwargs)
        Format a single message from the template.
    format_messages(**kwargs)
        Format multiple messages from the template.

    Notes
    -----
    This is an abstract base class that must be subclassed to provide
    concrete implementations. Subclasses should implement the `format_message`
    and `format_messages` methods according to their specific template
    formatting requirements.

    Examples
    --------
    >>> class MyTemplate(BasePromptTemplate):
    ...     def format_message(self, role=Role.USER, **kwargs):
    ...         # Implementation here
    ...         pass
    ...     
    ...     def format_messages(self, **kwargs):
    ...         # Implementation here
    ...         pass
    >>> 
    >>> template = MyTemplate(template_str="Hello {name}!")
    >>> message = template.format_message(name="World")
    """

    template_str: str

    def format_message(self, role: Union[Role, str] = Role.USER, **kwargs) -> Message:
        """
        Format a single message from the template.

        Parameters
        ----------
        role : Union[Role, str], default=Role.USER
            The role of the message (e.g., 'user', 'assistant', 'system').
        **kwargs
            Additional keyword arguments to be substituted into the template.

        Returns
        -------
        Message
            A formatted message object.

        Raises
        ------
        NotImplementedError
            This method must be implemented by subclasses.
        """
        raise NotImplementedError(f"format_message is not implemented in class {self.__class__.__name__}")

    def format_messages(self, **kwargs) -> List[Message]:
        """
        Format multiple messages from the template.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments to be substituted into the template.

        Returns
        -------
        List[Message]
            A list of formatted message objects.

        Raises
        ------
        NotImplementedError
            This method must be implemented by subclasses.
        """
        raise NotImplementedError(f"format_messages is not implemented in class {self.__class__.__name__}")

class FstringPromptTemplate(BasePromptTemplate):
    """    
    This template implementation uses Python's f-string syntax (braces `{}`).

    Methods
    -------
    format_message(role, **kwargs)
        Format a single message from the template.

    Notes
    -----
    This template supports single message rendering via `format_message()`.
    The template uses Python's built-in `str.format()` method for variable
    substitution, which provides basic formatting capabilities.

    Examples
    --------
    Basic usage:
    >>> template = FstringPromptTemplate("Hello {name}, you are {age} years old.")
    >>> message = template.format_message(role="user", name="Alice", age=25)

    With context:
    >>> template = FstringPromptTemplate('''
    ... Context: {context}
    ... Question: {question}
    ... Please provide a helpful answer.
    ... ''')
    >>> message = template.format_message(
    ...     role="system", 
    ...     context="Python programming", 
    ...     question="What is a decorator?"
    ... )

    Multiple variables:
    >>> template = FstringPromptTemplate("{greeting} {name}! Today is {date}.")
    >>> message = template.format_message(
    ...     role="assistant",
    ...     greeting="Good morning",
    ...     name="Bob", 
    ...     date="Monday"
    ... )
    """

    def format_message(self, role: Union[Role, str], **kwargs) -> Message:
        """
        Format a single message from the template.

        Parameters
        ----------
        role : Union[Role, str]
            The role of the message (e.g., 'user', 'assistant', 'system').
            Required parameter for this template implementation.
        **kwargs
            Keyword arguments containing values for all variables referenced
            in the template string. All variables must be provided.

        Returns
        -------
        Message
            A formatted message object with the specified role and rendered content.

        Raises
        ------
        PromptRenderError
            If any variables referenced in the template are missing from
            the provided keyword arguments.
        """
        if isinstance(role, str):
            role = Role(role)

        all_vars = self._find_variables()
        missing_vars = set(all_vars) - set(kwargs.keys())
        if missing_vars:
            raise PromptRenderError(f"Missing variables that are required to render the prompt template: {', '.join(missing_vars)}")

        rendered = self.template_str.format(**kwargs)
        return Message.from_text(text=rendered, role=role)

    def _find_variables(self) -> List[str]:
        """
        Extract variable names from the template string.

        Returns
        -------
        List[str]
            A list of unique variable names found in the template string,
            in the order they first appear. Variable names are extracted
            from curly brace syntax `{variable_name}`.
        """
        var_list = re.findall(r'{([^}]+)}', self.template_str)
        var_list = [var.strip() for var in var_list]
        return unique_list_in_order(var_list)

class EjinjaPromptTemplate(BasePromptTemplate):
    """
    Extended Jinja2-based prompt template with custom message blocks.

    This template implementation extends the standard Jinja2 syntax with custom
    `{% msg %}` blocks to create structured Message objects. It supports both
    single message and multiple message rendering with variable substitution
    and content block parsing.

    Attributes
    ----------
    _env_template : Template
        The compiled Jinja2 template object.
    _render_cache : MemoryCache
        Cache for rendered template results to improve performance.

    Methods
    -------
    format_message(role, **kwargs)
        Format a single message from the template.
    format_messages(**kwargs)
        Format multiple messages from the template.

    Notes
    -----
    This template supports two rendering modes:

    1. **Single Message Mode**: Use `format_message()` to render one message.    
    2. **Multiple Messages Mode**: Use `format_messages()` to render multiple messages.

    Examples
    --------
    Single message with role in template:
    >>> template = EjinjaPromptTemplate('''
    ... {% msg role="system" %}
    ... You are a helpful assistant. User name: {{ name }}
    ... {% endmsg %}
    ... ''')
    >>> message = template.format_message(name="Alice")

    Single message with role as parameter:
    >>> template = EjinjaPromptTemplate("Hello {{ name }}, how are you?")
    >>> message = template.format_message(role="user", name="Bob")

    Multiple messages:
    >>> template = EjinjaPromptTemplate('''
    ... {% msg role="system" %}You are helpful{% endmsg %}
    ... {% msg role="user" %}Hello {{ name }}{% endmsg %}
    ... ''')
    >>> messages = template.format_messages(name="Charlie")
    """

    _env_template: Template
    _render_cache: MemoryCache

    def __init__(self, template_str: str):
        """
        Initialize the EjinjaPromptTemplate.

        Parameters
        ----------
        template_str : str
            The Jinja2 template string with optional `{% msg %}` blocks.
        """
        super().__init__(template_str=template_str)
        self._env_template = env.from_string(template_str)
        self._render_cache = MemoryCache()

    def format_message(self, role: Union[Role, str] = None, **kwargs) -> Message:
        """
        Format a single message from the template.

        Parameters
        ----------
        role : Union[Role, str], optional
            The role of the message. If the template contains a `{% msg %}` block,
            this parameter should be None as the role will be extracted from
            the template. If no `{% msg %}` block exists, this parameter is required.
        **kwargs
            Additional keyword arguments to be substituted into the template.

        Returns
        -------
        Message
            A formatted message object with the specified role and content.

        Raises
        ------
        PromptSyntaxError
            If the template contains more than one `{% msg %}` block.
        PromptRenderError
            If role parameter conflicts with template-defined role, or if
            no role is specified when template has no `{% msg %}` block.
        """
        if isinstance(role, str):
            role = Role(role)

        rendered = self._env_template.render(**kwargs)
        match_list = re.findall(r"{%\s*msg\s*role=\"(.*?)\"\s*%}(.*?){%\s*endmsg\s*%}", rendered)
        if len(match_list) > 1:
            raise PromptSyntaxError(
                f"It is required to just have one {{% msg %}} block in the template, "
                f"but got {len(match_list)}"
            )
        elif len(match_list) == 1:
            if role is not None:
                raise PromptRenderError(
                    f"If you want to render a single message, the role has to be only specified in the template "
                    f"and not be passed as an argument to the \"format_message\" method in {type(self).__name__}"
                )
            role, content = match_list[0][0], match_list[0][1]
        else:
            if role is None:
                raise PromptRenderError(
                    f"If you want to render a template without {{% msg %}} blocks, the role has to be specified "
                    f"as an argument to the \"format_message\" method in {type(self).__name__}"
                )
            role, content = role, rendered
        return Message.from_text(text=content, role=role)

    def format_messages(self, **kwargs) -> List[Message]:
        """
        Format multiple messages from the template.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments to be substituted into the template.

        Returns
        -------
        List[Message]
            A list of formatted message objects. Each line of the rendered
            template should be a valid JSON representation of a Message object.
            If no valid messages are found but content exists, a default user
            message is created.

        Raises
        ------
        PromptRenderError
            If any line in the rendered template is not a valid JSON
            representation of a Message object.

        Notes
        -----
        This method uses caching to improve performance for repeated calls
        with the same parameters. The rendered template is cached based on
        the provided keyword arguments.
        """
        rendered = self._render_cache.get(kwargs)
        if not rendered:
            rendered = self._env_template.render(kwargs)
            self._render_cache.set(kwargs, rendered)

        messages: List[Message] = []
        for line in rendered.strip().split("\n"):
            try:
                messages.append(Message.model_validate_json(line))
            except Exception:
                raise PromptRenderError(
                    f"It is required to wrap each content in a {{% msg %}} block when calling the "
                    f"\"format_messages\" method of {type(self).__name__}, but got: {line}"
                )

        if not messages and rendered.strip():
            messages.append(_chat_message_from_text(role="user", content=rendered))
        return messages