coco-framework 0.1.0__py3-none-any.whl

coco_core/prompts/few_shot_with_templates.py

@@ -0,0 +1,239 @@
+"""Prompt template that contains few shot examples."""
+
+from pathlib import Path
+from typing import Any
+
+from pydantic import ConfigDict, model_validator
+from typing_extensions import Self
+
+from coco_core._api import deprecated
+from coco_core.example_selectors import BaseExampleSelector
+from coco_core.prompts.prompt import PromptTemplate
+from coco_core.prompts.string import (
+    DEFAULT_FORMATTER_MAPPING,
+    PromptTemplateFormat,
+    StringPromptTemplate,
+)
+
+
+class FewShotPromptWithTemplates(StringPromptTemplate):
+    """Prompt template that contains few shot examples."""
+
+    examples: list[dict[str, Any]] | None = None
+    """Examples to format into the prompt.
+
+    Either this or `example_selector` should be provided.
+    """
+
+    example_selector: BaseExampleSelector | None = None
+    """`ExampleSelector` to choose the examples to format into the prompt.
+
+    Either this or `examples` should be provided.
+    """
+
+    example_prompt: PromptTemplate
+    """`PromptTemplate` used to format an individual example."""
+
+    suffix: StringPromptTemplate
+    """A `PromptTemplate` to put after the examples."""
+
+    example_separator: str = "\n\n"
+    """String separator used to join the prefix, the examples, and suffix."""
+
+    prefix: StringPromptTemplate | None = None
+    """A `PromptTemplate` to put before the examples."""
+
+    template_format: PromptTemplateFormat = "f-string"
+    """The format of the prompt template.
+
+    Options are: `'f-string'`, `'jinja2'`, `'mustache'`.
+    """
+
+    validate_template: bool = False
+    """Whether or not to try validating the template."""
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "prompts", "few_shot_with_templates"]`
+        """
+        return ["coco", "prompts", "few_shot_with_templates"]
+
+    @model_validator(mode="before")
+    @classmethod
+    def check_examples_and_selector(cls, values: dict[str, Any]) -> Any:
+        """Check that one and only one of examples/example_selector are provided."""
+        examples = values.get("examples")
+        example_selector = values.get("example_selector")
+        if examples and example_selector:
+            msg = "Only one of 'examples' and 'example_selector' should be provided"
+            raise ValueError(msg)
+
+        if examples is None and example_selector is None:
+            msg = "One of 'examples' and 'example_selector' should be provided"
+            raise ValueError(msg)
+
+        return values
+
+    @model_validator(mode="after")
+    def template_is_valid(self) -> Self:
+        """Check that prefix, suffix, and input variables are consistent."""
+        if self.validate_template:
+            input_variables = self.input_variables
+            expected_input_variables = set(self.suffix.input_variables)
+            expected_input_variables |= set(self.partial_variables)
+            if self.prefix is not None:
+                expected_input_variables |= set(self.prefix.input_variables)
+            missing_vars = expected_input_variables.difference(input_variables)
+            if missing_vars:
+                msg = (
+                    f"Got input_variables={input_variables}, but based on "
+                    f"prefix/suffix expected {expected_input_variables}"
+                )
+                raise ValueError(msg)
+        else:
+            self.input_variables = sorted(
+                set(self.suffix.input_variables)
+                | set(self.prefix.input_variables if self.prefix else [])
+                - set(self.partial_variables)
+            )
+        return self
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
+
+    def _get_examples(self, **kwargs: Any) -> list[dict[str, Any]]:
+        if self.examples is not None:
+            return self.examples
+        if self.example_selector is not None:
+            return self.example_selector.select_examples(kwargs)
+        msg = "One of 'examples' and 'example_selector' should be provided"
+        raise ValueError(msg)
+
+    async def _aget_examples(self, **kwargs: Any) -> list[dict[str, Any]]:
+        if self.examples is not None:
+            return self.examples
+        if self.example_selector is not None:
+            return await self.example_selector.aselect_examples(kwargs)
+        msg = "One of 'examples' and 'example_selector' should be provided"
+        raise ValueError(msg)
+
+    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.
+
+        Example:
+            ```python
+            prompt.format(variable1="foo")
+            ```
+        """
+        kwargs = self._merge_partial_and_user_variables(**kwargs)
+        # Get the examples to use.
+        examples = self._get_examples(**kwargs)
+        # Format the examples.
+        example_strings = [
+            self.example_prompt.format(**example) for example in examples
+        ]
+        # Create the overall prefix.
+        if self.prefix is None:
+            prefix = ""
+        else:
+            prefix_kwargs = {
+                k: v for k, v in kwargs.items() if k in self.prefix.input_variables
+            }
+            for k in prefix_kwargs:
+                kwargs.pop(k)
+            prefix = self.prefix.format(**prefix_kwargs)
+
+        # Create the overall suffix
+        suffix_kwargs = {
+            k: v for k, v in kwargs.items() if k in self.suffix.input_variables
+        }
+        for k in suffix_kwargs:
+            kwargs.pop(k)
+        suffix = self.suffix.format(
+            **suffix_kwargs,
+        )
+
+        pieces = [prefix, *example_strings, suffix]
+        template = self.example_separator.join([piece for piece in pieces if piece])
+        # Format the template with the input variables.
+        return DEFAULT_FORMATTER_MAPPING[self.template_format](template, **kwargs)
+
+    async def aformat(self, **kwargs: Any) -> str:
+        """Async 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)
+        # Get the examples to use.
+        examples = await self._aget_examples(**kwargs)
+        # Format the examples.
+        example_strings = [
+            # We can use the sync method here as PromptTemplate doesn't block
+            self.example_prompt.format(**example)
+            for example in examples
+        ]
+        # Create the overall prefix.
+        if self.prefix is None:
+            prefix = ""
+        else:
+            prefix_kwargs = {
+                k: v for k, v in kwargs.items() if k in self.prefix.input_variables
+            }
+            for k in prefix_kwargs:
+                kwargs.pop(k)
+            prefix = await self.prefix.aformat(**prefix_kwargs)
+
+        # Create the overall suffix
+        suffix_kwargs = {
+            k: v for k, v in kwargs.items() if k in self.suffix.input_variables
+        }
+        for k in suffix_kwargs:
+            kwargs.pop(k)
+        suffix = await self.suffix.aformat(
+            **suffix_kwargs,
+        )
+
+        pieces = [prefix, *example_strings, suffix]
+        template = self.example_separator.join([piece for piece in pieces if piece])
+        # Format the template with the input variables.
+        return DEFAULT_FORMATTER_MAPPING[self.template_format](template, **kwargs)
+
+    @property
+    def _prompt_type(self) -> str:
+        """Return the prompt type key."""
+        return "few_shot_with_templates"
+
+    @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 save(self, file_path: Path | str) -> None:
+        """Save the prompt to a file.
+
+        Args:
+            file_path: The path to save the prompt to.
+
+        Raises:
+            ValueError: If `example_selector` is provided.
+        """
+        if self.example_selector:
+            msg = "Saving an example selector is not currently supported"
+            raise ValueError(msg)
+        return super().save(file_path)

coco_core/prompts/image.py

@@ -0,0 +1,177 @@
+"""Image prompt template for a multimodal model."""
+
+from typing import Any, Literal, cast
+
+from pydantic import Field
+
+from coco_core.prompt_values import ImagePromptValue, ImageURL, PromptValue
+from coco_core.prompts.base import BasePromptTemplate
+from coco_core.prompts.string import (
+    DEFAULT_FORMATTER_MAPPING,
+    PromptTemplateFormat,
+    get_template_variables,
+)
+from coco_core.runnables import run_in_executor
+
+
+class ImagePromptTemplate(BasePromptTemplate[ImageURL]):
+    """Image prompt template for a multimodal model.
+
+    Example:
+        ```python
+        prompt = ImagePromptTemplate(
+            input_variables=["image_id"],
+            template={"url": "https://example.com/{image_id}.png", "detail": "high"},
+            template_format="f-string",
+        )
+        prompt.format(image_id="cat")
+        # {"url": "https://example.com/cat.png", "detail": "high"}
+        ```
+    """
+
+    template: dict[str, Any] = Field(default_factory=dict)
+    """Template for the prompt."""
+
+    template_format: PromptTemplateFormat = "f-string"
+    """The format of the prompt template.
+
+    Options are: `'f-string'`, `'mustache'`, `'jinja2'`.
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Create an image prompt template.
+
+        Raises:
+            ValueError: If the input variables contain `'url'`, `'path'`, or
+                `'detail'`.
+        """
+        if "input_variables" not in kwargs:
+            kwargs["input_variables"] = []
+
+        overlap = set(kwargs["input_variables"]) & {"url", "path", "detail"}
+        if overlap:
+            msg = (
+                "input_variables for the image template cannot contain"
+                " any of 'url', 'path', or 'detail'."
+                f" Found: {overlap}"
+            )
+            raise ValueError(msg)
+
+        template = kwargs.get("template", {})
+        template_format = kwargs.get("template_format", "f-string")
+        for value in template.values():
+            if isinstance(value, str):
+                get_template_variables(value, template_format)
+
+        super().__init__(**kwargs)
+
+    @property
+    def _prompt_type(self) -> str:
+        """Return the prompt type key."""
+        return "image-prompt"
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "prompts", "image"]`
+        """
+        return ["coco", "prompts", "image"]
+
+    def format_prompt(self, **kwargs: Any) -> PromptValue:
+        """Format the prompt with the inputs.
+
+        Args:
+            **kwargs: Any arguments to be passed to the prompt template.
+
+        Returns:
+            A formatted string.
+        """
+        return ImagePromptValue(image_url=self.format(**kwargs))
+
+    async def aformat_prompt(self, **kwargs: Any) -> PromptValue:
+        """Async format the prompt with the inputs.
+
+        Args:
+            **kwargs: Any arguments to be passed to the prompt template.
+
+        Returns:
+            A formatted string.
+        """
+        return ImagePromptValue(image_url=await self.aformat(**kwargs))
+
+    def format(
+        self,
+        **kwargs: Any,
+    ) -> ImageURL:
+        """Format the prompt with the inputs.
+
+        Args:
+            **kwargs: Any arguments to be passed to the prompt template.
+
+        Returns:
+            A formatted string.
+
+        Raises:
+            ValueError: If the url is not provided.
+            ValueError: If the url is not a string.
+            ValueError: If `'path'` is provided in the template or kwargs.
+
+        Example:
+            ```python
+            prompt.format(variable1="foo")
+            ```
+        """
+        formatted = {}
+        for k, v in self.template.items():
+            if isinstance(v, str):
+                formatted[k] = DEFAULT_FORMATTER_MAPPING[self.template_format](
+                    v, **kwargs
+                )
+            else:
+                formatted[k] = v
+        url = kwargs.get("url") or formatted.get("url")
+        if kwargs.get("path") or formatted.get("path"):
+            msg = (
+                "Loading images from 'path' has been removed as of 0.3.15 for security "
+                "reasons. Please specify images by 'url'."
+            )
+            raise ValueError(msg)
+        detail = kwargs.get("detail") or formatted.get("detail")
+        if not url:
+            msg = "Must provide url."
+            raise ValueError(msg)
+        if not isinstance(url, str):
+            msg = "url must be a string."
+            raise ValueError(msg)  # noqa: TRY004
+        output: ImageURL = {"url": url}
+        if detail:
+            # Don't check literal values here: let the API check them
+            output["detail"] = cast("Literal['auto', 'low', 'high']", detail)
+        return output
+
+    async def aformat(self, **kwargs: Any) -> ImageURL:
+        """Async format the prompt with the inputs.
+
+        Args:
+            **kwargs: Any arguments to be passed to the prompt template.
+
+        Returns:
+            A formatted string.
+        """
+        return await run_in_executor(None, self.format, **kwargs)
+
+    def pretty_repr(
+        self,
+        html: bool = False,  # noqa: FBT001,FBT002
+    ) -> str:
+        """Return a pretty representation of the prompt.
+
+        Args:
+            html: Whether to return an html formatted string.
+
+        Returns:
+            A pretty representation of the prompt.
+        """
+        raise NotImplementedError