codekingpro/portable-devtools

1"""`BasePrompt` schema definition."""2 3from __future__ import annotations4 5import warnings6from abc import ABC, abstractmethod7from string import Formatter8from typing import TYPE_CHECKING, Any, Literal, cast9 10from pydantic import BaseModel, create_model11from typing_extensions import override12 13from langchain_core.prompt_values import PromptValue, StringPromptValue14from langchain_core.prompts.base import BasePromptTemplate15from langchain_core.utils import get_colored_text, mustache16from langchain_core.utils.formatting import formatter17from langchain_core.utils.interactive_env import is_interactive_env18 19if TYPE_CHECKING:20    from collections.abc import Callable, Sequence21 22try:23    from jinja2 import meta24    from jinja2.sandbox import SandboxedEnvironment25 26    _HAS_JINJA2 = True27except ImportError:28    _HAS_JINJA2 = False29 30PromptTemplateFormat = Literal["f-string", "mustache", "jinja2"]31 32 33def jinja2_formatter(template: str, /, **kwargs: Any) -> str:34    """Format a template using jinja2.35 36    !!! warning "Security"37 38        As of LangChain 0.0.329, this method uses Jinja2's `SandboxedEnvironment` by39        default. However, this sandboxing should be treated as a best-effort approach40        rather than a guarantee of security.41 42        Do not accept jinja2 templates from untrusted sources as they may lead43        to arbitrary Python code execution.44 45        [More information.](https://jinja.palletsprojects.com/en/3.1.x/sandbox/)46 47    Args:48        template: The template string.49        **kwargs: The variables to format the template with.50 51    Returns:52        The formatted string.53 54    Raises:55        ImportError: If jinja2 is not installed.56    """57    if not _HAS_JINJA2:58        msg = (59            "jinja2 not installed, which is needed to use the jinja2_formatter. "60            "Please install it with `pip install jinja2`."61            "Please be cautious when using jinja2 templates. "62            "Do not expand jinja2 templates using unverified or user-controlled "63            "inputs as that can result in arbitrary Python code execution."64        )65        raise ImportError(msg)66 67    # Use Jinja2's SandboxedEnvironment which blocks access to dunder attributes68    # (e.g., __class__, __globals__) to prevent sandbox escapes.69    # Note: regular attribute access (e.g., {{obj.attr}}) and method calls are70    # still allowed. This is a best-effort measure — do not use with untrusted71    # templates.72    return SandboxedEnvironment().from_string(template).render(**kwargs)73 74 75def validate_jinja2(template: str, input_variables: list[str]) -> None:76    """Validate that the input variables are valid for the template.77 78    Issues a warning if missing or extra variables are found.79 80    Args:81        template: The template string.82        input_variables: The input variables.83    """84    input_variables_set = set(input_variables)85    valid_variables = _get_jinja2_variables_from_template(template)86    missing_variables = valid_variables - input_variables_set87    extra_variables = input_variables_set - valid_variables88 89    warning_message = ""90    if missing_variables:91        warning_message += f"Missing variables: {missing_variables} "92 93    if extra_variables:94        warning_message += f"Extra variables: {extra_variables}"95 96    if warning_message:97        warnings.warn(warning_message.strip(), stacklevel=7)98 99 100def _get_jinja2_variables_from_template(template: str) -> set[str]:101    if not _HAS_JINJA2:102        msg = (103            "jinja2 not installed, which is needed to use the jinja2_formatter. "104            "Please install it with `pip install jinja2`."105        )106        raise ImportError(msg)107    env = SandboxedEnvironment()108    ast = env.parse(template)109    return meta.find_undeclared_variables(ast)110 111 112def mustache_formatter(template: str, /, **kwargs: Any) -> str:113    """Format a template using mustache.114 115    Args:116        template: The template string.117        **kwargs: The variables to format the template with.118 119    Returns:120        The formatted string.121    """122    return mustache.render(template, kwargs)123 124 125def mustache_template_vars(126    template: str,127) -> set[str]:128    """Get the top-level variables from a mustache template.129 130    For nested variables like `{{person.name}}`, only the top-level key (`person`) is131    returned.132 133    Args:134        template: The template string.135 136    Returns:137        The top-level variables from the template.138    """139    variables: set[str] = set()140    section_depth = 0141    for type_, key in mustache.tokenize(template):142        if type_ == "end":143            section_depth -= 1144        elif (145            type_ in {"variable", "section", "inverted section", "no escape"}146            and key != "."147            and section_depth == 0148        ):149            variables.add(key.split(".")[0])150        if type_ in {"section", "inverted section"}:151            section_depth += 1152    return variables153 154 155Defs = dict[str, "Defs"]156 157 158def mustache_schema(template: str) -> type[BaseModel]:159    """Get the variables from a mustache template.160 161    Args:162        template: The template string.163 164    Returns:165        The variables from the template as a Pydantic model.166    """167    fields = {}168    prefix: tuple[str, ...] = ()169    section_stack: list[tuple[str, ...]] = []170    for type_, key in mustache.tokenize(template):171        if key == ".":172            continue173        if type_ == "end":174            if section_stack:175                prefix = section_stack.pop()176        elif type_ in {"section", "inverted section"}:177            section_stack.append(prefix)178            prefix += tuple(key.split("."))179            fields[prefix] = False180        elif type_ in {"variable", "no escape"}:181            fields[prefix + tuple(key.split("."))] = True182 183    for fkey, fval in fields.items():184        fields[fkey] = fval and not any(185            is_subsequence(fkey, k) for k in fields if k != fkey186        )187    defs: Defs = {}  # None means leaf node188    while fields:189        field, is_leaf = fields.popitem()190        current = defs191        for part in field[:-1]:192            current = current.setdefault(part, {})193        current.setdefault(field[-1], "" if is_leaf else {})  # type: ignore[arg-type]194    return _create_model_recursive("PromptInput", defs)195 196 197def _create_model_recursive(name: str, defs: Defs) -> type[BaseModel]:198    return cast(199        "type[BaseModel]",200        create_model(  # type: ignore[call-overload]201            name,202            **{203                k: (_create_model_recursive(k, v), None) if v else (type(v), None)204                for k, v in defs.items()205            },206        ),207    )208 209 210DEFAULT_FORMATTER_MAPPING: dict[str, Callable[..., str]] = {211    "f-string": formatter.format,212    "mustache": mustache_formatter,213    "jinja2": jinja2_formatter,214}215 216DEFAULT_VALIDATOR_MAPPING: dict[str, Callable] = {217    "f-string": formatter.validate_input_variables,218    "jinja2": validate_jinja2,219}220 221 222def _parse_f_string_fields(template: str) -> list[tuple[str, str | None]]:223    fields: list[tuple[str, str | None]] = []224    for _, field_name, format_spec, _ in Formatter().parse(template):225        if field_name is not None:226            fields.append((field_name, format_spec))227    return fields228 229 230def validate_f_string_template(template: str) -> list[str]:231    """Validate an f-string template and return its input variables."""232    input_variables = set()233    for var, format_spec in _parse_f_string_fields(template):234        if "." in var or "[" in var or "]" in var:235            msg = (236                f"Invalid variable name {var!r} in f-string template. "237                f"Variable names cannot contain attribute "238                f"access (.) or indexing ([])."239            )240            raise ValueError(msg)241 242        if var.isdigit():243            msg = (244                f"Invalid variable name {var!r} in f-string template. "245                f"Variable names cannot be all digits as they are interpreted "246                f"as positional arguments."247            )248            raise ValueError(msg)249 250        if format_spec and ("{" in format_spec or "}" in format_spec):251            msg = (252                "Invalid format specifier in f-string template. "253                "Nested replacement fields are not allowed."254            )255            raise ValueError(msg)256 257        input_variables.add(var)258 259    return sorted(input_variables)260 261 262def check_valid_template(263    template: str, template_format: str, input_variables: list[str]264) -> None:265    """Check that template string is valid.266 267    Args:268        template: The template string.269        template_format: The template format.270 271            Should be one of `'f-string'` or `'jinja2'`.272        input_variables: The input variables.273 274    Raises:275        ValueError: If the template format is not supported.276        ValueError: If the prompt schema is invalid.277    """278    try:279        validator_func = DEFAULT_VALIDATOR_MAPPING[template_format]280    except KeyError as exc:281        msg = (282            f"Invalid template format {template_format!r}, should be one of"283            f" {list(DEFAULT_FORMATTER_MAPPING)}."284        )285        raise ValueError(msg) from exc286    if template_format == "f-string":287        validate_f_string_template(template)288    try:289        validator_func(template, input_variables)290    except (KeyError, IndexError) as exc:291        msg = (292            "Invalid prompt schema; check for mismatched or missing input parameters"293            f" from {input_variables}."294        )295        raise ValueError(msg) from exc296 297 298def get_template_variables(template: str, template_format: str) -> list[str]:299    """Get the variables from the template.300 301    Args:302        template: The template string.303        template_format: The template format.304 305            Should be one of `'f-string'`, `'mustache'` or `'jinja2'`.306 307    Returns:308        The variables from the template.309 310    Raises:311        ValueError: If the template format is not supported.312    """313    input_variables: list[str] | set[str]314    if template_format == "jinja2":315        # Get the variables for the template316        input_variables = sorted(_get_jinja2_variables_from_template(template))317    elif template_format == "f-string":318        input_variables = validate_f_string_template(template)319    elif template_format == "mustache":320        input_variables = mustache_template_vars(template)321    else:322        msg = f"Unsupported template format: {template_format}"323        raise ValueError(msg)324 325    return sorted(input_variables)326 327 328class StringPromptTemplate(BasePromptTemplate, ABC):329    """String prompt that exposes the format method, returning a prompt."""330 331    @classmethod332    def get_lc_namespace(cls) -> list[str]:333        """Get the namespace of the LangChain object.334 335        Returns:336            `["langchain", "prompts", "base"]`337        """338        return ["langchain", "prompts", "base"]339 340    def format_prompt(self, **kwargs: Any) -> PromptValue:341        """Format the prompt with the inputs.342 343        Args:344            **kwargs: Any arguments to be passed to the prompt template.345 346        Returns:347            A formatted string.348        """349        return StringPromptValue(text=self.format(**kwargs))350 351    async def aformat_prompt(self, **kwargs: Any) -> PromptValue:352        """Async format the prompt with the inputs.353 354        Args:355            **kwargs: Any arguments to be passed to the prompt template.356 357        Returns:358            A formatted string.359        """360        return StringPromptValue(text=await self.aformat(**kwargs))361 362    @override363    @abstractmethod364    def format(self, **kwargs: Any) -> str: ...365 366    def pretty_repr(367        self,368        html: bool = False,  # noqa: FBT001,FBT002369    ) -> str:370        """Get a pretty representation of the prompt.371 372        Args:373            html: Whether to return an HTML-formatted string.374 375        Returns:376            A pretty representation of the prompt.377        """378        # TODO: handle partials379        dummy_vars = {380            input_var: "{" + f"{input_var}" + "}" for input_var in self.input_variables381        }382        if html:383            dummy_vars = {384                k: get_colored_text(v, "yellow") for k, v in dummy_vars.items()385            }386        return self.format(**dummy_vars)387 388    def pretty_print(self) -> None:389        """Print a pretty representation of the prompt."""390        print(self.pretty_repr(html=is_interactive_env()))  # noqa: T201391 392 393def is_subsequence(child: Sequence, parent: Sequence) -> bool:394    """Return `True` if child is subsequence of parent."""395    if len(child) == 0 or len(parent) == 0:396        return False397    if len(parent) < len(child):398        return False399    return all(child[i] == parent[i] for i in range(len(child)))400