src/vsp/llm/prompt_text.py

from typing import Any


class _DefaultDict(dict[str, Any]):
    def __missing__(self, key: str) -> str:
        return "{" + key + "}"


class PromptText:
    """
    A class for managing and formatting prompt templates with dynamic inputs.

    This class allows for the creation of prompt templates with placeholders
    that can be filled with dynamic content. It provides methods to update
    inputs and retrieve the formatted prompt.
    """

    def __init__(self, prompt_template: str, inputs: dict[str, Any] | None = None) -> None:
        """
        Initialize a PromptText object.

        Args:
            prompt_template (str): The template string for the prompt.
            inputs (dict[str, Any] | None): Initial input values to substitute into the template.
                If None, an empty dictionary will be used.

        Raises:
            ValueError: If the template is invalid (e.g., unmatched braces).
        """

        self.prompt_template = prompt_template
        self.inputs = inputs or {}
        self._prompt = ""
        self._update_prompt()

    def _update_prompt(self) -> None:
        """
        Update the prompt by formatting the template with the current inputs.

        This method is called internally whenever the inputs are updated.

        Raises:
            ValueError: If the template is invalid (e.g., unmatched braces).
        """
        default_inputs = _DefaultDict(self.inputs)
        try:
            self._prompt = self.prompt_template.format_map(default_inputs)
        except ValueError as e:
            raise ValueError(f"Invalid prompt template: {e}")

    def upsert_inputs(self, new_inputs: dict[str, Any]) -> None:
        """
        Update the inputs and regenerate the prompt.

        This method updates the existing inputs with new values and then
        regenerates the prompt using the updated inputs.

        Args:
            new_inputs (dict[str, Any]): A dictionary of new input values to update.

        Raises:
            ValueError: If the template becomes invalid after updating inputs.
        """

        self.inputs.update(new_inputs)
        self._update_prompt()

    def get_prompt(self) -> str:
        """
        Retrieve the current formatted prompt.

        Returns:
            str: The current prompt string after formatting with inputs.
        """
        return self._prompt