coco-framework 0.1.0__py3-none-any.whl

coco_core/prompts/loading.py

@@ -0,0 +1,296 @@
+"""Load prompts."""
+
+import json
+import logging
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from coco_core._api import deprecated
+from coco_core.output_parsers.string import StrOutputParser
+from coco_core.prompts.base import BasePromptTemplate
+from coco_core.prompts.chat import ChatPromptTemplate
+from coco_core.prompts.few_shot import FewShotPromptTemplate
+from coco_core.prompts.prompt import PromptTemplate
+
+URL_BASE = "https://raw.githubusercontent.com/hwchase17/langchain-hub/master/prompts/"
+logger = logging.getLogger(__name__)
+
+
+def _validate_path(path: Path) -> None:
+    """Reject absolute paths and `..` traversal components.
+
+    Args:
+        path: The path to validate.
+
+    Raises:
+        ValueError: If the path is absolute or contains `..` components.
+    """
+    if path.is_absolute():
+        msg = (
+            f"Path '{path}' is absolute. Absolute paths are not allowed "
+            f"when loading prompt configurations to prevent path traversal "
+            f"attacks. Use relative paths instead, or pass "
+            f"`allow_dangerous_paths=True` if you trust the input."
+        )
+        raise ValueError(msg)
+    if ".." in path.parts:
+        msg = (
+            f"Path '{path}' contains '..' components. Directory traversal "
+            f"sequences are not allowed when loading prompt configurations. "
+            f"Use direct relative paths instead, or pass "
+            f"`allow_dangerous_paths=True` if you trust the input."
+        )
+        raise ValueError(msg)
+
+
+@deprecated(
+    since="1.2.21",
+    removal="2.0.0",
+    alternative="Use `dumpd`/`dumps` from `coco_core.load` to serialize "
+    "prompts and `load`/`loads` to deserialize them.",
+)
+def load_prompt_from_config(
+    config: dict[str, Any], *, allow_dangerous_paths: bool = False
+) -> BasePromptTemplate[str]:
+    """Load prompt from config dict.
+
+    Args:
+        config: Dict containing the prompt configuration.
+        allow_dangerous_paths: If `False` (default), file paths in the
+            config (such as `template_path`, `examples`, and
+            `example_prompt_path`) are validated to reject absolute paths
+            and directory traversal (`..`) sequences. Set to `True` only
+            if you trust the source of the config.
+
+    Returns:
+        A `PromptTemplate` object.
+
+    Raises:
+        ValueError: If the prompt type is not supported.
+    """
+    if "_type" not in config:
+        logger.warning("No `_type` key found, defaulting to `prompt`.")
+    config_type = config.pop("_type", "prompt")
+
+    if config_type not in type_to_loader_dict:
+        msg = f"Loading {config_type} prompt not supported"
+        raise ValueError(msg)
+
+    prompt_loader = type_to_loader_dict[config_type]
+    return prompt_loader(config, allow_dangerous_paths=allow_dangerous_paths)
+
+
+def _load_template(
+    var_name: str, config: dict[str, Any], *, allow_dangerous_paths: bool = False
+) -> dict[str, Any]:
+    """Load template from the path if applicable."""
+    # Check if template_path exists in config.
+    if f"{var_name}_path" in config:
+        # If it does, make sure template variable doesn't also exist.
+        if var_name in config:
+            msg = f"Both `{var_name}_path` and `{var_name}` cannot be provided."
+            raise ValueError(msg)
+        # Pop the template path from the config.
+        template_path = Path(config.pop(f"{var_name}_path"))
+        if not allow_dangerous_paths:
+            _validate_path(template_path)
+        # Resolve symlinks before checking the suffix so that a symlink named
+        # "exploit.txt" pointing to a non-.txt file is caught.
+        resolved_path = template_path.resolve()
+        # Load the template.
+        if resolved_path.suffix == ".txt":
+            template = resolved_path.read_text(encoding="utf-8")
+        else:
+            msg = (
+                f"Unsupported template file format: '{resolved_path.suffix}'. "
+                "Only '.txt' files are supported."
+            )
+            raise ValueError(msg)
+        # Set the template variable to the extracted variable.
+        config[var_name] = template
+    return config
+
+
+def _load_examples(
+    config: dict[str, Any], *, allow_dangerous_paths: bool = False
+) -> dict[str, Any]:
+    """Load examples if necessary."""
+    if isinstance(config["examples"], list):
+        pass
+    elif isinstance(config["examples"], str):
+        path = Path(config["examples"])
+        if not allow_dangerous_paths:
+            _validate_path(path)
+        with path.open(encoding="utf-8") as f:
+            if path.suffix == ".json":
+                examples = json.load(f)
+            elif path.suffix in {".yaml", ".yml"}:
+                examples = yaml.safe_load(f)
+            else:
+                msg = "Invalid file format. Only json or yaml formats are supported."
+                raise ValueError(msg)
+        config["examples"] = examples
+    else:
+        msg = "Invalid examples format. Only list or string are supported."
+        raise ValueError(msg)  # noqa:TRY004
+    return config
+
+
+def _load_output_parser(config: dict[str, Any]) -> dict[str, Any]:
+    """Load output parser."""
+    if config_ := config.get("output_parser"):
+        if output_parser_type := config_.get("_type") != "default":
+            msg = f"Unsupported output parser {output_parser_type}"
+            raise ValueError(msg)
+        config["output_parser"] = StrOutputParser(**config_)
+    return config
+
+
+def _load_few_shot_prompt(
+    config: dict[str, Any], *, allow_dangerous_paths: bool = False
+) -> FewShotPromptTemplate:
+    """Load the "few shot" prompt from the config."""
+    # Load the suffix and prefix templates.
+    config = _load_template(
+        "suffix", config, allow_dangerous_paths=allow_dangerous_paths
+    )
+    config = _load_template(
+        "prefix", config, allow_dangerous_paths=allow_dangerous_paths
+    )
+    # Load the example prompt.
+    if "example_prompt_path" in config:
+        if "example_prompt" in config:
+            msg = (
+                "Only one of example_prompt and example_prompt_path should "
+                "be specified."
+            )
+            raise ValueError(msg)
+        example_prompt_path = Path(config.pop("example_prompt_path"))
+        if not allow_dangerous_paths:
+            _validate_path(example_prompt_path)
+        config["example_prompt"] = load_prompt(
+            example_prompt_path, allow_dangerous_paths=allow_dangerous_paths
+        )
+    else:
+        config["example_prompt"] = load_prompt_from_config(
+            config["example_prompt"], allow_dangerous_paths=allow_dangerous_paths
+        )
+    # Load the examples.
+    config = _load_examples(config, allow_dangerous_paths=allow_dangerous_paths)
+    config = _load_output_parser(config)
+    return FewShotPromptTemplate(**config)
+
+
+def _load_prompt(
+    config: dict[str, Any], *, allow_dangerous_paths: bool = False
+) -> PromptTemplate:
+    """Load the prompt template from config."""
+    # Load the template from disk if necessary.
+    config = _load_template(
+        "template", config, allow_dangerous_paths=allow_dangerous_paths
+    )
+    config = _load_output_parser(config)
+
+    template_format = config.get("template_format", "f-string")
+    if template_format == "jinja2":
+        # Disabled due to:
+        # https://github.com/langchain-ai/langchain/issues/4394
+        msg = (
+            f"Loading templates with '{template_format}' format is no longer supported "
+            f"since it can lead to arbitrary code execution. Please migrate to using "
+            f"the 'f-string' template format, which does not suffer from this issue."
+        )
+        raise ValueError(msg)
+
+    return PromptTemplate(**config)
+
+
+@deprecated(
+    since="1.2.21",
+    removal="2.0.0",
+    alternative="Use `dumpd`/`dumps` from `coco_core.load` to serialize "
+    "prompts and `load`/`loads` to deserialize them.",
+)
+def load_prompt(
+    path: str | Path,
+    encoding: str | None = None,
+    *,
+    allow_dangerous_paths: bool = False,
+) -> BasePromptTemplate[str]:
+    """Unified method for loading a prompt from CocoHub or local filesystem.
+
+    Args:
+        path: Path to the prompt file.
+        encoding: Encoding of the file.
+        allow_dangerous_paths: If `False` (default), file paths referenced
+            inside the loaded config (such as `template_path`, `examples`,
+            and `example_prompt_path`) are validated to reject absolute paths
+            and directory traversal (`..`) sequences. Set to `True` only
+            if you trust the source of the config.
+
+    Returns:
+        A `PromptTemplate` object.
+
+    Raises:
+        RuntimeError: If the path is a CocoHub path.
+    """
+    if isinstance(path, str) and path.startswith("lc://"):
+        msg = (
+            "Loading from the deprecated github-based Hub is no longer supported. "
+            "Please use the new Coco Hub at https://smith.langchain.com/hub "
+            "instead."
+        )
+        raise RuntimeError(msg)
+    return _load_prompt_from_file(
+        path, encoding, allow_dangerous_paths=allow_dangerous_paths
+    )
+
+
+def _load_prompt_from_file(
+    file: str | Path,
+    encoding: str | None = None,
+    *,
+    allow_dangerous_paths: bool = False,
+) -> BasePromptTemplate[str]:
+    """Load prompt from file."""
+    # Convert file to a Path object.
+    file_path = Path(file)
+    # Load from either json or yaml.
+    if file_path.suffix == ".json":
+        with file_path.open(encoding=encoding) as f:
+            config = json.load(f)
+    elif file_path.suffix.endswith((".yaml", ".yml")):
+        with file_path.open(encoding=encoding) as f:
+            config = yaml.safe_load(f)
+    else:
+        msg = f"Got unsupported file type {file_path.suffix}"
+        raise ValueError(msg)
+    # Load the prompt from the config now.
+    return load_prompt_from_config(config, allow_dangerous_paths=allow_dangerous_paths)
+
+
+def _load_chat_prompt(
+    config: dict[str, Any],
+    *,
+    allow_dangerous_paths: bool = False,  # noqa: ARG001
+) -> ChatPromptTemplate:
+    """Load chat prompt from config."""
+    messages = config.pop("messages")
+    template = messages[0]["prompt"].pop("template") if messages else None
+    config.pop("input_variables")
+
+    if not template:
+        msg = "Can't load chat prompt without template"
+        raise ValueError(msg)
+
+    return ChatPromptTemplate.from_template(template=template, **config)
+
+
+type_to_loader_dict: dict[str, Callable[..., BasePromptTemplate[str]]] = {
+    "prompt": _load_prompt,
+    "few_shot": _load_few_shot_prompt,
+    "chat": _load_chat_prompt,
+}

coco_core/prompts/message.py

@@ -0,0 +1,97 @@
+"""Message prompt templates."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from coco_core.load import Serializable
+from coco_core.utils.interactive_env import is_interactive_env
+
+if TYPE_CHECKING:
+    from coco_core.messages import BaseMessage
+    from coco_core.prompts.chat import ChatPromptTemplate
+
+
+class BaseMessagePromptTemplate(Serializable, ABC):
+    """Base class for message prompt templates."""
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return `True` as this class is serializable."""
+        return True
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "prompts", "chat"]`
+        """
+        return ["coco", "prompts", "chat"]
+
+    @abstractmethod
+    def format_messages(self, **kwargs: Any) -> list[BaseMessage]:
+        """Format messages from kwargs.
+
+        Should return a list of `BaseMessage` objects.
+
+        Args:
+            **kwargs: Keyword arguments to use for formatting.
+
+        Returns:
+            List of `BaseMessage` objects.
+        """
+
+    async def aformat_messages(self, **kwargs: Any) -> list[BaseMessage]:
+        """Async format messages from kwargs.
+
+        Args:
+            **kwargs: Keyword arguments to use for formatting.
+
+        Returns:
+            List of `BaseMessage` objects.
+        """
+        return self.format_messages(**kwargs)
+
+    @property
+    @abstractmethod
+    def input_variables(self) -> list[str]:
+        """Input variables for this prompt template.
+
+        Returns:
+            List of input variables.
+        """
+
+    def pretty_repr(
+        self,
+        html: bool = False,  # noqa: FBT001,FBT002
+    ) -> str:
+        """Human-readable representation.
+
+        Args:
+            html: Whether to format as HTML.
+
+        Returns:
+            Human-readable representation.
+        """
+        raise NotImplementedError
+
+    def pretty_print(self) -> None:
+        """Print a human-readable representation."""
+        print(self.pretty_repr(html=is_interactive_env()))  # noqa: T201
+
+    def __add__(self, other: Any) -> ChatPromptTemplate:
+        """Combine two prompt templates.
+
+        Args:
+            other: Another prompt template.
+
+        Returns:
+            Combined prompt template.
+        """
+        # Import locally to avoid circular import.
+        from coco_core.prompts.chat import ChatPromptTemplate  # noqa: PLC0415
+
+        prompt = ChatPromptTemplate(messages=[self])
+        return prompt.__add__(other)

coco_core/prompts/prompt.py

@@ -0,0 +1,312 @@
+"""Prompt schema definition."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, model_validator
+from typing_extensions import override
+
+from coco_core.prompts.string import (
+    DEFAULT_FORMATTER_MAPPING,
+    PromptTemplateFormat,
+    StringPromptTemplate,
+    check_valid_template,
+    get_template_variables,
+    mustache_schema,
+)
+
+if TYPE_CHECKING:
+    from coco_core.runnables.config import RunnableConfig
+
+
+class PromptTemplate(StringPromptTemplate):
+    """Prompt template for a language model.
+
+    A prompt template consists of a string template. It accepts a set of parameters
+    from the user that can be used to generate a prompt for a language model.
+
+    The template can be formatted using either f-strings (default), jinja2, or mustache
+    syntax.
+
+    !!! warning "Security"
+
+        Prefer using `template_format='f-string'` instead of `template_format='jinja2'`,
+        or make sure to NEVER accept jinja2 templates from untrusted sources as they may
+        lead to arbitrary Python code execution.
+
+        As of Coco 0.0.329, Jinja2 templates will be rendered using Jinja2's
+        SandboxedEnvironment by default. This sand-boxing should be treated as a
+        best-effort approach rather than a guarantee of security, as it is an opt-out
+        rather than opt-in approach.
+
+        Despite the sandboxing, we recommend to never use jinja2 templates from
+        untrusted sources.
+
+    Example:
+        ```python
+        from coco_core.prompts import PromptTemplate
+
+        # Instantiation using from_template (recommended)
+        prompt = PromptTemplate.from_template("Say {foo}")
+        prompt.format(foo="bar")
+
+        # Instantiation using initializer
+        prompt = PromptTemplate(template="Say {foo}")
+        ```
+    """
+
+    @property
+    @override
+    def lc_attributes(self) -> dict[str, Any]:
+        return {
+            "template_format": self.template_format,
+        }
+
+    @classmethod
+    @override
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "prompts", "prompt"]`
+        """
+        return ["coco", "prompts", "prompt"]
+
+    template: str
+    """The prompt template."""
+
+    template_format: PromptTemplateFormat = "f-string"
+    """The format of the prompt template.
+
+    Options are: `'f-string'`, `'mustache'`, `'jinja2'`.
+    """
+
+    validate_template: bool = False
+    """Whether or not to try validating the template."""
+
+    @model_validator(mode="before")
+    @classmethod
+    def pre_init_validation(cls, values: dict[str, Any]) -> Any:
+        """Check that template and input variables are consistent."""
+        if values.get("template") is None:
+            # Will let pydantic fail with a ValidationError if template
+            # is not provided.
+            return values
+
+        # Set some default values based on the field defaults
+        values.setdefault("template_format", "f-string")
+        values.setdefault("partial_variables", {})
+
+        if values.get("validate_template"):
+            if values["template_format"] == "mustache":
+                msg = "Mustache templates cannot be validated."
+                raise ValueError(msg)
+
+            if "input_variables" not in values:
+                msg = "Input variables must be provided to validate the template."
+                raise ValueError(msg)
+
+            all_inputs = values["input_variables"] + list(values["partial_variables"])
+            check_valid_template(
+                values["template"], values["template_format"], all_inputs
+            )
+
+        if values["template_format"]:
+            values["input_variables"] = [
+                var
+                for var in get_template_variables(
+                    values["template"], values["template_format"]
+                )
+                if var not in values["partial_variables"]
+            ]
+
+        return values
+
+    @override
+    def get_input_schema(self, config: RunnableConfig | None = None) -> type[BaseModel]:
+        """Get the input schema for the prompt.
+
+        Args:
+            config: The runnable configuration.
+
+        Returns:
+            The input schema for the prompt.
+        """
+        if self.template_format != "mustache":
+            return super().get_input_schema(config)
+
+        return mustache_schema(self.template)
+
+    def __add__(self, other: Any) -> PromptTemplate:
+        """Override the `+` operator to allow for combining prompt templates.
+
+        Raises:
+            ValueError: If the template formats are not f-string or if there are
+                conflicting partial variables.
+            NotImplementedError: If the other object is not a `PromptTemplate` or str.
+
+        Returns:
+            A new `PromptTemplate` that is the combination of the two.
+        """
+        # Allow for easy combining
+        if isinstance(other, PromptTemplate):
+            if self.template_format != other.template_format:
+                msg = "Cannot add templates of different formats"
+                raise ValueError(msg)
+            input_variables = list(
+                set(self.input_variables) | set(other.input_variables)
+            )
+            template = self.template + other.template
+            # If any do not want to validate, then don't
+            validate_template = self.validate_template and other.validate_template
+            partial_variables = dict(self.partial_variables.items())
+            for k, v in other.partial_variables.items():
+                if k in partial_variables:
+                    msg = "Cannot have same variable partialed twice."
+                    raise ValueError(msg)
+                partial_variables[k] = v
+            return PromptTemplate(
+                template=template,
+                input_variables=input_variables,
+                partial_variables=partial_variables,
+                template_format=self.template_format,
+                validate_template=validate_template,
+            )
+        if isinstance(other, str):
+            prompt = PromptTemplate.from_template(
+                other,
+                template_format=self.template_format,
+            )
+            return self + prompt
+        msg = f"Unsupported operand type for +: {type(other)}"
+        raise NotImplementedError(msg)
+
+    @property
+    def _prompt_type(self) -> str:
+        """Return the prompt type key."""
+        return "prompt"
+
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt with the inputs.
+
+        Args:
+            **kwargs: Any arguments to be passed to the prompt template.
+
+        Returns:
+            A formatted string.
+        """
+        kwargs = self._merge_partial_and_user_variables(**kwargs)
+        return DEFAULT_FORMATTER_MAPPING[self.template_format](self.template, **kwargs)
+
+    @classmethod
+    def from_examples(
+        cls,
+        examples: list[str],
+        suffix: str,
+        input_variables: list[str],
+        example_separator: str = "\n\n",
+        prefix: str = "",
+        **kwargs: Any,
+    ) -> PromptTemplate:
+        """Take examples in list format with prefix and suffix to create a prompt.
+
+        Intended to be used as a way to dynamically create a prompt from examples.
+
+        Args:
+            examples: List of examples to use in the prompt.
+            suffix: String to go after the list of examples.
+
+                Should generally set up the user's input.
+            input_variables: A list of variable names the final prompt template will
+                expect.
+            example_separator: The separator to use in between examples.
+            prefix: String that should go before any examples.
+
+                Generally includes examples.
+
+        Returns:
+            The final prompt generated.
+        """
+        template = example_separator.join([prefix, *examples, suffix])
+        return cls(input_variables=input_variables, template=template, **kwargs)
+
+    @classmethod
+    def from_file(
+        cls,
+        template_file: str | Path,
+        encoding: str | None = None,
+        **kwargs: Any,
+    ) -> PromptTemplate:
+        """Load a prompt from a file.
+
+        Args:
+            template_file: The path to the file containing the prompt template.
+            encoding: The encoding system for opening the template file.
+
+                If not provided, will use the OS default.
+
+        Returns:
+            The prompt loaded from the file.
+        """
+        template = Path(template_file).read_text(encoding=encoding)
+        return cls.from_template(template=template, **kwargs)
+
+    @classmethod
+    def from_template(
+        cls,
+        template: str,
+        *,
+        template_format: PromptTemplateFormat = "f-string",
+        partial_variables: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> PromptTemplate:
+        """Load a prompt template from a template.
+
+        !!! warning "Security"
+
+            Prefer using `template_format='f-string'` instead of
+            `template_format='jinja2'`, or make sure to NEVER accept jinja2 templates
+            from untrusted sources as they may lead to arbitrary Python code execution.
+
+            As of Coco 0.0.329, Jinja2 templates will be rendered using Jinja2's
+            SandboxedEnvironment by default. This sand-boxing should be treated as a
+            best-effort approach rather than a guarantee of security, as it is an
+            opt-out rather than opt-in approach.
+
+            Despite the sandboxing, we recommend to never use jinja2 templates from
+            untrusted sources.
+
+        Args:
+            template: The template to load.
+            template_format: The format of the template.
+
+                Use `jinja2` for jinja2, `mustache` for mustache, and `f-string` for
+                f-strings.
+            partial_variables: A dictionary of variables that can be used to partially
+                fill in the template.
+
+                For example, if the template is `'{variable1} {variable2}'`, and
+                `partial_variables` is `{"variable1": "foo"}`, then the final prompt
+                will be `'foo {variable2}'`.
+            **kwargs: Any other arguments to pass to the prompt template.
+
+        Returns:
+            The prompt template loaded from the template.
+        """
+        input_variables = get_template_variables(template, template_format)
+        partial_variables_ = partial_variables or {}
+
+        if partial_variables_:
+            input_variables = [
+                var for var in input_variables if var not in partial_variables_
+            ]
+
+        return cls(
+            input_variables=input_variables,
+            template=template,
+            template_format=template_format,
+            partial_variables=partial_variables_,
+            **kwargs,
+        )