coco-framework 0.1.0__py3-none-any.whl

coco_core/prompts/dict.py

@@ -0,0 +1,177 @@
+"""Dictionary prompt template."""
+
+import warnings
+from functools import cached_property
+from typing import Any, Literal, cast
+
+from pydantic import model_validator
+from typing_extensions import override
+
+from coco_core.load import dumpd
+from coco_core.prompts.string import (
+    DEFAULT_FORMATTER_MAPPING,
+    get_template_variables,
+)
+from coco_core.runnables import RunnableConfig, RunnableSerializable
+from coco_core.runnables.config import ensure_config
+
+
+class DictPromptTemplate(RunnableSerializable[dict[str, Any], dict[str, Any]]):
+    """Template represented by a dictionary.
+
+    Recognizes variables in f-string or mustache formatted string dict values.
+
+    Does NOT recognize variables in dict keys. Applies recursively.
+
+    Example:
+        ```python
+        prompt = DictPromptTemplate(
+            template={
+                "type": "text",
+                "text": "Hello {name}",
+                "metadata": {"source": "{source}"},
+            },
+            template_format="f-string",
+        )
+        prompt.format(name="Alice", source="docs")
+        # {
+        #     "type": "text",
+        #     "text": "Hello Alice",
+        #     "metadata": {"source": "docs"},
+        # }
+        ```
+    """
+
+    template: dict[str, Any]
+    template_format: Literal["f-string", "mustache"]
+
+    @model_validator(mode="after")
+    def validate_template(self) -> "DictPromptTemplate":
+        """Validate that the template structure contains only safe variables."""
+        _get_input_variables(self.template, self.template_format)
+        return self
+
+    @property
+    def input_variables(self) -> list[str]:
+        """Template input variables."""
+        return _get_input_variables(self.template, self.template_format)
+
+    def format(self, **kwargs: Any) -> dict[str, Any]:
+        """Format the prompt with the inputs.
+
+        Returns:
+            A formatted dict.
+        """
+        return _insert_input_variables(self.template, kwargs, self.template_format)
+
+    async def aformat(self, **kwargs: Any) -> dict[str, Any]:
+        """Format the prompt with the inputs.
+
+        Returns:
+            A formatted dict.
+        """
+        return self.format(**kwargs)
+
+    @override
+    def invoke(
+        self, input: dict[str, Any], config: RunnableConfig | None = None, **kwargs: Any
+    ) -> dict[str, Any]:
+        return self._call_with_config(
+            lambda x: self.format(**x),
+            input,
+            ensure_config(config),
+            run_type="prompt",
+            serialized=self._serialized,
+            **kwargs,
+        )
+
+    @property
+    def _prompt_type(self) -> str:
+        return "dict-prompt"
+
+    @cached_property
+    def _serialized(self) -> dict[str, Any]:
+        # self is always a Serializable object in this case, thus the result is
+        # guaranteed to be a dict since dumpd uses the default callback, which uses
+        # obj.to_json which always returns TypedDict subclasses
+        return cast("dict[str, Any]", dumpd(self))
+
+    @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_core", "prompts", "dict"]`
+        """
+        return ["coco_core", "prompts", "dict"]
+
+    def pretty_repr(self, *, html: bool = False) -> str:
+        """Human-readable representation.
+
+        Args:
+            html: Whether to format as HTML.
+
+        Returns:
+            Human-readable representation.
+        """
+        raise NotImplementedError
+
+
+def _get_input_variables(
+    template: dict[str, Any], template_format: Literal["f-string", "mustache"]
+) -> list[str]:
+    input_variables = []
+    for v in template.values():
+        if isinstance(v, str):
+            input_variables += get_template_variables(v, template_format)
+        elif isinstance(v, dict):
+            input_variables += _get_input_variables(v, template_format)
+        elif isinstance(v, (list, tuple)):
+            for x in v:
+                if isinstance(x, str):
+                    input_variables += get_template_variables(x, template_format)
+                elif isinstance(x, dict):
+                    input_variables += _get_input_variables(x, template_format)
+    return list(set(input_variables))
+
+
+def _insert_input_variables(
+    template: dict[str, Any],
+    inputs: dict[str, Any],
+    template_format: Literal["f-string", "mustache"],
+) -> dict[str, Any]:
+    formatted: dict[str, Any] = {}
+    formatter = DEFAULT_FORMATTER_MAPPING[template_format]
+    for k, v in template.items():
+        if isinstance(v, str):
+            formatted[k] = formatter(v, **inputs)
+        elif isinstance(v, dict):
+            if k == "image_url" and "path" in v:
+                msg = (
+                    "Specifying image inputs via file path in environments with "
+                    "user-input paths is a security vulnerability. Out of an abundance "
+                    "of caution, the utility has been removed to prevent possible "
+                    "misuse."
+                )
+                warnings.warn(msg, stacklevel=2)
+            formatted[k] = _insert_input_variables(v, inputs, template_format)
+        elif isinstance(v, (list, tuple)):
+            formatted_v: list[Any] = []
+            for x in v:
+                if isinstance(x, str):
+                    formatted_v.append(formatter(x, **inputs))
+                elif isinstance(x, dict):
+                    formatted_v.append(
+                        _insert_input_variables(x, inputs, template_format)
+                    )
+                else:
+                    formatted_v.append(x)
+            formatted[k] = type(v)(formatted_v)
+        else:
+            formatted[k] = v
+    return formatted

coco_core/prompts/few_shot.py

@@ -0,0 +1,483 @@
+"""Prompt template that contains few shot examples."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    model_validator,
+)
+from typing_extensions import override
+
+from coco_core._api import deprecated
+from coco_core.example_selectors import BaseExampleSelector
+from coco_core.messages import BaseMessage, get_buffer_string
+from coco_core.prompts.chat import BaseChatPromptTemplate
+from coco_core.prompts.message import BaseMessagePromptTemplate
+from coco_core.prompts.prompt import PromptTemplate
+from coco_core.prompts.string import (
+    DEFAULT_FORMATTER_MAPPING,
+    StringPromptTemplate,
+    check_valid_template,
+    get_template_variables,
+)
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from typing_extensions import Self
+
+
+class _FewShotPromptTemplateMixin(BaseModel):
+    """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.
+    """
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
+
+    @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.
+
+        Args:
+            values: The values to check.
+
+        Returns:
+            The values if they are valid.
+
+        Raises:
+            ValueError: If neither or both `examples` and `example_selector` are
+                provided.
+            ValueError: If both `examples` and `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
+
+    def _get_examples(self, **kwargs: Any) -> list[dict[str, Any]]:
+        """Get the examples to use for formatting the prompt.
+
+        Args:
+            **kwargs: Keyword arguments to be passed to the example selector.
+
+        Returns:
+            List of examples.
+
+        Raises:
+            ValueError: If neither `examples` nor `example_selector` are provided.
+        """
+        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]]:
+        """Async get the examples to use for formatting the prompt.
+
+        Args:
+            **kwargs: Keyword arguments to be passed to the example selector.
+
+        Returns:
+            List of examples.
+
+        Raises:
+            ValueError: If neither `examples` nor `example_selector` are provided.
+        """
+        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)
+
+
+class FewShotPromptTemplate(_FewShotPromptTemplateMixin, StringPromptTemplate):
+    """Prompt template that contains few shot examples."""
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return `False` as this class is not serializable."""
+        return False
+
+    validate_template: bool = False
+    """Whether or not to try validating the template."""
+
+    example_prompt: PromptTemplate
+    """`PromptTemplate` used to format an individual example."""
+
+    suffix: str
+    """A prompt template string to put after the examples."""
+
+    example_separator: str = "\n\n"
+    """String separator used to join the prefix, the examples, and suffix."""
+
+    prefix: str = ""
+    """A prompt template string to put before the examples."""
+
+    template_format: Literal["f-string", "jinja2"] = "f-string"
+    """The format of the prompt template.
+
+    Options are: `'f-string'`, `'jinja2'`.
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the few shot prompt template."""
+        if "input_variables" not in kwargs and "example_prompt" in kwargs:
+            kwargs["input_variables"] = kwargs["example_prompt"].input_variables
+        super().__init__(**kwargs)
+
+    @model_validator(mode="after")
+    def template_is_valid(self) -> Self:
+        """Check that prefix, suffix, and input variables are consistent."""
+        if self.validate_template:
+            check_valid_template(
+                self.prefix + self.suffix,
+                self.template_format,
+                self.input_variables + list(self.partial_variables),
+            )
+        elif self.template_format:
+            self.input_variables = [
+                var
+                for var in get_template_variables(
+                    self.prefix + self.suffix, self.template_format
+                )
+                if var not in self.partial_variables
+            ]
+        return self
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
+
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt with inputs generating a string.
+
+        Use this method to generate a string representation of a prompt.
+
+        Args:
+            **kwargs: Keyword arguments to use for formatting.
+
+        Returns:
+            A string representation of the prompt.
+        """
+        kwargs = self._merge_partial_and_user_variables(**kwargs)
+        # Get the examples to use.
+        examples = self._get_examples(**kwargs)
+        examples = [
+            {k: e[k] for k in self.example_prompt.input_variables} for e in examples
+        ]
+        # Format the examples.
+        example_strings = [
+            self.example_prompt.format(**example) for example in examples
+        ]
+        # Create the overall template.
+        pieces = [self.prefix, *example_strings, self.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 inputs generating a string.
+
+        Use this method to generate a string representation of a prompt.
+
+        Args:
+            **kwargs: Keyword arguments to use for formatting.
+
+        Returns:
+            A string representation of the prompt.
+        """
+        kwargs = self._merge_partial_and_user_variables(**kwargs)
+        # Get the examples to use.
+        examples = await self._aget_examples(**kwargs)
+        examples = [
+            {k: e[k] for k in self.example_prompt.input_variables} for e in examples
+        ]
+        # Format the examples.
+        example_strings = [
+            await self.example_prompt.aformat(**example) for example in examples
+        ]
+        # Create the overall template.
+        pieces = [self.prefix, *example_strings, self.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"
+
+    @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 template to a file.
+
+        Args:
+            file_path: The path to save the prompt template 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)
+
+
+class FewShotChatMessagePromptTemplate(
+    BaseChatPromptTemplate, _FewShotPromptTemplateMixin
+):
+    """Chat prompt template that supports few-shot examples.
+
+    The high level structure of produced by this prompt template is a list of messages
+    consisting of prefix message(s), example message(s), and suffix message(s).
+
+    This structure enables creating a conversation with intermediate examples like:
+
+    ```txt
+    System: You are a helpful AI Assistant
+
+    Human: What is 2+2?
+
+    AI: 4
+
+    Human: What is 2+3?
+
+    AI: 5
+
+    Human: What is 4+4?
+    ```
+
+    This prompt template can be used to generate a fixed list of examples or else to
+    dynamically select examples based on the input.
+
+    Examples:
+        Prompt template with a fixed list of examples (matching the sample
+        conversation above):
+
+        ```python
+        from coco_core.prompts import (
+            FewShotChatMessagePromptTemplate,
+            ChatPromptTemplate,
+        )
+
+        examples = [
+            {"input": "2+2", "output": "4"},
+            {"input": "2+3", "output": "5"},
+        ]
+
+        example_prompt = ChatPromptTemplate.from_messages(
+            [
+                ("human", "What is {input}?"),
+                ("ai", "{output}"),
+            ]
+        )
+
+        few_shot_prompt = FewShotChatMessagePromptTemplate(
+            examples=examples,
+            # This is a prompt template used to format each individual example.
+            example_prompt=example_prompt,
+        )
+
+        final_prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", "You are a helpful AI Assistant"),
+                few_shot_prompt,
+                ("human", "{input}"),
+            ]
+        )
+        final_prompt.format(input="What is 4+4?")
+        ```
+
+        Prompt template with dynamically selected examples:
+
+        ```python
+        from coco_core.prompts import SemanticSimilarityExampleSelector
+        from coco_core.embeddings import OpenAIEmbeddings
+        from coco_core.vectorstores import Chroma
+
+        examples = [
+            {"input": "2+2", "output": "4"},
+            {"input": "2+3", "output": "5"},
+            {"input": "2+4", "output": "6"},
+            # ...
+        ]
+
+        to_vectorize = [" ".join(example.values()) for example in examples]
+        embeddings = OpenAIEmbeddings()
+        vectorstore = Chroma.from_texts(to_vectorize, embeddings, metadatas=examples)
+        example_selector = SemanticSimilarityExampleSelector(vectorstore=vectorstore)
+
+        from coco_core import SystemMessage
+        from coco_core.prompts import HumanMessagePromptTemplate
+        from coco_core.prompts.few_shot import FewShotChatMessagePromptTemplate
+
+        few_shot_prompt = FewShotChatMessagePromptTemplate(
+            # Which variable(s) will be passed to the example selector.
+            input_variables=["input"],
+            example_selector=example_selector,
+            # Define how each example will be formatted.
+            # In this case, each example will become 2 messages:
+            # 1 human, and 1 AI
+            example_prompt=(
+                HumanMessagePromptTemplate.from_template("{input}")
+                + AIMessagePromptTemplate.from_template("{output}")
+            ),
+        )
+        # Define the overall prompt.
+        final_prompt = (
+            SystemMessagePromptTemplate.from_template("You are a helpful AI Assistant")
+            + few_shot_prompt
+            + HumanMessagePromptTemplate.from_template("{input}")
+        )
+        # Show the prompt
+        print(final_prompt.format_messages(input="What's 3+3?"))  # noqa: T201
+
+        # Use within an LLM
+        from coco_anthropic import ChatAnthropic
+
+        chain = final_prompt | ChatAnthropic(model="claude-sonnet-5")
+        chain.invoke({"input": "What's 3+3?"})
+        ```
+    """
+
+    input_variables: list[str] = Field(default_factory=list)
+    """A list of the names of the variables the prompt template will use to pass to
+    the `example_selector`, if provided.
+    """
+
+    example_prompt: BaseMessagePromptTemplate | BaseChatPromptTemplate
+    """The class to format each example."""
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return `False` as this class is not serializable."""
+        return False
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+    )
+
+    def format_messages(self, **kwargs: Any) -> list[BaseMessage]:
+        """Format kwargs into a list of messages.
+
+        Args:
+            **kwargs: Keyword arguments to use for filling in templates in messages.
+
+        Returns:
+            A list of formatted messages with all template variables filled in.
+        """
+        # Get the examples to use.
+        examples = self._get_examples(**kwargs)
+        examples = [
+            {k: e[k] for k in self.example_prompt.input_variables} for e in examples
+        ]
+        # Format the examples.
+        return [
+            message
+            for example in examples
+            for message in self.example_prompt.format_messages(**example)
+        ]
+
+    async def aformat_messages(self, **kwargs: Any) -> list[BaseMessage]:
+        """Async format kwargs into a list of messages.
+
+        Args:
+            **kwargs: Keyword arguments to use for filling in templates in messages.
+
+        Returns:
+            A list of formatted messages with all template variables filled in.
+        """
+        # Get the examples to use.
+        examples = await self._aget_examples(**kwargs)
+        examples = [
+            {k: e[k] for k in self.example_prompt.input_variables} for e in examples
+        ]
+        # Format the examples.
+        return [
+            message
+            for example in examples
+            for message in await self.example_prompt.aformat_messages(**example)
+        ]
+
+    def format(self, **kwargs: Any) -> str:
+        """Format the prompt with inputs generating a string.
+
+        Use this method to generate a string representation of a prompt consisting of
+        chat messages.
+
+        Useful for feeding into a string-based completion language model or debugging.
+
+        Args:
+            **kwargs: Keyword arguments to use for formatting.
+
+        Returns:
+            A string representation of the prompt
+        """
+        messages = self.format_messages(**kwargs)
+        return get_buffer_string(messages)
+
+    async def aformat(self, **kwargs: Any) -> str:
+        """Async format the prompt with inputs generating a string.
+
+        Use this method to generate a string representation of a prompt consisting of
+        chat messages.
+
+        Useful for feeding into a string-based completion language model or debugging.
+
+        Args:
+            **kwargs: Keyword arguments to use for formatting.
+
+        Returns:
+            A string representation of the prompt
+        """
+        messages = await self.aformat_messages(**kwargs)
+        return get_buffer_string(messages)
+
+    @override
+    def pretty_repr(self, html: bool = False) -> str:
+        """Return a pretty representation of the prompt template.
+
+        Args:
+            html: Whether or not to return an HTML formatted string.
+
+        Returns:
+            A pretty representation of the prompt template.
+        """
+        raise NotImplementedError