coco-framework 0.1.0__py3-none-any.whl

coco_core/prompts/string.py

@@ -0,0 +1,399 @@
+"""`BasePrompt` schema definition."""
+
+from __future__ import annotations
+
+import warnings
+from abc import ABC, abstractmethod
+from string import Formatter
+from typing import TYPE_CHECKING, Any, Literal, cast
+
+from pydantic import BaseModel, create_model
+from typing_extensions import override
+
+from coco_core.prompt_values import PromptValue, StringPromptValue
+from coco_core.prompts.base import BasePromptTemplate
+from coco_core.utils import get_colored_text, mustache
+from coco_core.utils.formatting import formatter
+from coco_core.utils.interactive_env import is_interactive_env
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+try:
+    from jinja2 import meta
+    from jinja2.sandbox import SandboxedEnvironment
+
+    _HAS_JINJA2 = True
+except ImportError:
+    _HAS_JINJA2 = False
+
+PromptTemplateFormat = Literal["f-string", "mustache", "jinja2"]
+
+
+def jinja2_formatter(template: str, /, **kwargs: Any) -> str:
+    """Format a template using jinja2.
+
+    !!! warning "Security"
+
+        As of Coco 0.0.329, this method uses Jinja2's `SandboxedEnvironment` by
+        default. However, this sandboxing should be treated as a best-effort approach
+        rather than a guarantee of security.
+
+        Do not accept jinja2 templates from untrusted sources as they may lead
+        to arbitrary Python code execution.
+
+        [More information.](https://jinja.palletsprojects.com/en/3.1.x/sandbox/)
+
+    Args:
+        template: The template string.
+        **kwargs: The variables to format the template with.
+
+    Returns:
+        The formatted string.
+
+    Raises:
+        ImportError: If jinja2 is not installed.
+    """
+    if not _HAS_JINJA2:
+        msg = (
+            "jinja2 not installed, which is needed to use the jinja2_formatter. "
+            "Please install it with `pip install jinja2`."
+            "Please be cautious when using jinja2 templates. "
+            "Do not expand jinja2 templates using unverified or user-controlled "
+            "inputs as that can result in arbitrary Python code execution."
+        )
+        raise ImportError(msg)
+
+    # Use Jinja2's SandboxedEnvironment which blocks access to dunder attributes
+    # (e.g., __class__, __globals__) to prevent sandbox escapes.
+    # Note: regular attribute access (e.g., {{obj.attr}}) and method calls are
+    # still allowed. This is a best-effort measure — do not use with untrusted
+    # templates.
+    return SandboxedEnvironment().from_string(template).render(**kwargs)
+
+
+def validate_jinja2(template: str, input_variables: list[str]) -> None:
+    """Validate that the input variables are valid for the template.
+
+    Issues a warning if missing or extra variables are found.
+
+    Args:
+        template: The template string.
+        input_variables: The input variables.
+    """
+    input_variables_set = set(input_variables)
+    valid_variables = _get_jinja2_variables_from_template(template)
+    missing_variables = valid_variables - input_variables_set
+    extra_variables = input_variables_set - valid_variables
+
+    warning_message = ""
+    if missing_variables:
+        warning_message += f"Missing variables: {missing_variables} "
+
+    if extra_variables:
+        warning_message += f"Extra variables: {extra_variables}"
+
+    if warning_message:
+        warnings.warn(warning_message.strip(), stacklevel=7)
+
+
+def _get_jinja2_variables_from_template(template: str) -> set[str]:
+    if not _HAS_JINJA2:
+        msg = (
+            "jinja2 not installed, which is needed to use the jinja2_formatter. "
+            "Please install it with `pip install jinja2`."
+        )
+        raise ImportError(msg)
+    env = SandboxedEnvironment()
+    ast = env.parse(template)
+    return meta.find_undeclared_variables(ast)
+
+
+def mustache_formatter(template: str, /, **kwargs: Any) -> str:
+    """Format a template using mustache.
+
+    Args:
+        template: The template string.
+        **kwargs: The variables to format the template with.
+
+    Returns:
+        The formatted string.
+    """
+    return mustache.render(template, kwargs)
+
+
+def mustache_template_vars(
+    template: str,
+) -> set[str]:
+    """Get the top-level variables from a mustache template.
+
+    For nested variables like `{{person.name}}`, only the top-level key (`person`) is
+    returned.
+
+    Args:
+        template: The template string.
+
+    Returns:
+        The top-level variables from the template.
+    """
+    variables: set[str] = set()
+    section_depth = 0
+    for type_, key in mustache.tokenize(template):
+        if type_ == "end":
+            section_depth -= 1
+        elif (
+            type_ in {"variable", "section", "inverted section", "no escape"}
+            and key != "."
+            and section_depth == 0
+        ):
+            variables.add(key.split(".")[0])
+        if type_ in {"section", "inverted section"}:
+            section_depth += 1
+    return variables
+
+
+Defs = dict[str, "Defs"]
+
+
+def mustache_schema(template: str) -> type[BaseModel]:
+    """Get the variables from a mustache template.
+
+    Args:
+        template: The template string.
+
+    Returns:
+        The variables from the template as a Pydantic model.
+    """
+    fields = {}
+    prefix: tuple[str, ...] = ()
+    section_stack: list[tuple[str, ...]] = []
+    for type_, key in mustache.tokenize(template):
+        if key == ".":
+            continue
+        if type_ == "end":
+            if section_stack:
+                prefix = section_stack.pop()
+        elif type_ in {"section", "inverted section"}:
+            section_stack.append(prefix)
+            prefix += tuple(key.split("."))
+            fields[prefix] = False
+        elif type_ in {"variable", "no escape"}:
+            fields[prefix + tuple(key.split("."))] = True
+
+    for fkey, fval in fields.items():
+        fields[fkey] = fval and not any(
+            is_subsequence(fkey, k) for k in fields if k != fkey
+        )
+    defs: Defs = {}  # None means leaf node
+    while fields:
+        field, is_leaf = fields.popitem()
+        current = defs
+        for part in field[:-1]:
+            current = current.setdefault(part, {})
+        current.setdefault(field[-1], "" if is_leaf else {})  # type: ignore[arg-type]
+    return _create_model_recursive("PromptInput", defs)
+
+
+def _create_model_recursive(name: str, defs: Defs) -> type[BaseModel]:
+    return cast(
+        "type[BaseModel]",
+        create_model(  # type: ignore[call-overload]
+            name,
+            **{
+                k: (_create_model_recursive(k, v), None) if v else (type(v), None)
+                for k, v in defs.items()
+            },
+        ),
+    )
+
+
+DEFAULT_FORMATTER_MAPPING: dict[str, Callable[..., str]] = {
+    "f-string": formatter.format,
+    "mustache": mustache_formatter,
+    "jinja2": jinja2_formatter,
+}
+
+DEFAULT_VALIDATOR_MAPPING: dict[str, Callable[[str, list[str]], None]] = {
+    "f-string": formatter.validate_input_variables,
+    "jinja2": validate_jinja2,
+}
+
+
+def _parse_f_string_fields(template: str) -> list[tuple[str, str | None]]:
+    fields: list[tuple[str, str | None]] = []
+    for _, field_name, format_spec, _ in Formatter().parse(template):
+        if field_name is not None:
+            fields.append((field_name, format_spec))
+    return fields
+
+
+def validate_f_string_template(template: str) -> list[str]:
+    """Validate an f-string template and return its input variables."""
+    input_variables = set()
+    for var, format_spec in _parse_f_string_fields(template):
+        if "." in var or "[" in var or "]" in var:
+            msg = (
+                f"Invalid variable name {var!r} in f-string template. "
+                f"Variable names cannot contain attribute "
+                f"access (.) or indexing ([])."
+            )
+            raise ValueError(msg)
+
+        if var.isdigit():
+            msg = (
+                f"Invalid variable name {var!r} in f-string template. "
+                f"Variable names cannot be all digits as they are interpreted "
+                f"as positional arguments."
+            )
+            raise ValueError(msg)
+
+        if format_spec and ("{" in format_spec or "}" in format_spec):
+            msg = (
+                "Invalid format specifier in f-string template. "
+                "Nested replacement fields are not allowed."
+            )
+            raise ValueError(msg)
+
+        input_variables.add(var)
+
+    return sorted(input_variables)
+
+
+def check_valid_template(
+    template: str, template_format: str, input_variables: list[str]
+) -> None:
+    """Check that template string is valid.
+
+    Args:
+        template: The template string.
+        template_format: The template format.
+
+            Should be one of `'f-string'` or `'jinja2'`.
+        input_variables: The input variables.
+
+    Raises:
+        ValueError: If the template format is not supported.
+        ValueError: If the prompt schema is invalid.
+    """
+    try:
+        validator_func = DEFAULT_VALIDATOR_MAPPING[template_format]
+    except KeyError as exc:
+        msg = (
+            f"Invalid template format {template_format!r}, should be one of"
+            f" {list(DEFAULT_FORMATTER_MAPPING)}."
+        )
+        raise ValueError(msg) from exc
+    if template_format == "f-string":
+        validate_f_string_template(template)
+    try:
+        validator_func(template, input_variables)
+    except (KeyError, IndexError) as exc:
+        msg = (
+            "Invalid prompt schema; check for mismatched or missing input parameters"
+            f" from {input_variables}."
+        )
+        raise ValueError(msg) from exc
+
+
+def get_template_variables(template: str, template_format: str) -> list[str]:
+    """Get the variables from the template.
+
+    Args:
+        template: The template string.
+        template_format: The template format.
+
+            Should be one of `'f-string'`, `'mustache'` or `'jinja2'`.
+
+    Returns:
+        The variables from the template.
+
+    Raises:
+        ValueError: If the template format is not supported.
+    """
+    input_variables: list[str] | set[str]
+    if template_format == "jinja2":
+        # Get the variables for the template
+        input_variables = sorted(_get_jinja2_variables_from_template(template))
+    elif template_format == "f-string":
+        input_variables = validate_f_string_template(template)
+    elif template_format == "mustache":
+        input_variables = mustache_template_vars(template)
+    else:
+        msg = f"Unsupported template format: {template_format}"
+        raise ValueError(msg)
+
+    return sorted(input_variables)
+
+
+class StringPromptTemplate(BasePromptTemplate[str], ABC):
+    """String prompt that exposes the format method, returning a prompt."""
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "prompts", "base"]`
+        """
+        return ["coco", "prompts", "base"]
+
+    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 StringPromptValue(text=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 StringPromptValue(text=await self.aformat(**kwargs))
+
+    @override
+    @abstractmethod
+    def format(self, **kwargs: Any) -> str: ...
+
+    def pretty_repr(
+        self,
+        html: bool = False,  # noqa: FBT001,FBT002
+    ) -> str:
+        """Get a pretty representation of the prompt.
+
+        Args:
+            html: Whether to return an HTML-formatted string.
+
+        Returns:
+            A pretty representation of the prompt.
+        """
+        # TODO: handle partials
+        dummy_vars = {
+            input_var: "{" + f"{input_var}" + "}" for input_var in self.input_variables
+        }
+        if html:
+            dummy_vars = {
+                k: get_colored_text(v, "yellow") for k, v in dummy_vars.items()
+            }
+        return self.format(**dummy_vars)
+
+    def pretty_print(self) -> None:
+        """Print a pretty representation of the prompt."""
+        print(self.pretty_repr(html=is_interactive_env()))  # noqa: T201
+
+
+def is_subsequence(child: Sequence[Any], parent: Sequence[Any]) -> bool:
+    """Return `True` if child is subsequence of parent."""
+    if len(child) == 0 or len(parent) == 0:
+        return False
+    if len(parent) < len(child):
+        return False
+    return all(child[i] == parent[i] for i in range(len(child)))

coco_core/prompts/structured.py

@@ -0,0 +1,214 @@
+"""Structured prompt template for a language model."""
+
+from collections.abc import (
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Iterator,
+    Mapping,
+    Sequence,
+)
+from typing import (
+    Any,
+    overload,
+)
+
+from pydantic import BaseModel, Field
+from typing_extensions import override
+
+from coco_core._api.beta_decorator import beta
+from coco_core.language_models.base import BaseLanguageModel
+from coco_core.prompt_values import PromptValue
+from coco_core.prompts.chat import (
+    ChatPromptTemplate,
+    MessageLikeRepresentation,
+)
+from coco_core.prompts.string import PromptTemplateFormat
+from coco_core.runnables.base import (
+    Other,
+    Runnable,
+    RunnableSequence,
+    RunnableSerializable,
+)
+from coco_core.utils import get_pydantic_field_names
+
+
+@beta()
+class StructuredPrompt(ChatPromptTemplate):
+    """Structured prompt template for a language model."""
+
+    schema_: dict[str, Any] | type
+    """Schema for the structured prompt."""
+
+    structured_output_kwargs: dict[str, Any] = Field(default_factory=dict)
+
+    def __init__(
+        self,
+        messages: Sequence[MessageLikeRepresentation],
+        schema_: dict[str, Any] | type[BaseModel] | None = None,
+        *,
+        structured_output_kwargs: dict[str, Any] | None = None,
+        template_format: PromptTemplateFormat = "f-string",
+        **kwargs: Any,
+    ) -> None:
+        """Create a structured prompt template.
+
+        Args:
+            messages: Sequence of messages.
+            schema_: Schema for the structured prompt.
+            structured_output_kwargs: Additional kwargs for structured output.
+            template_format: Template format for the prompt.
+
+        Raises:
+            ValueError: If schema is not provided.
+        """
+        schema_ = schema_ or kwargs.pop("schema", None)
+        if not schema_:
+            err_msg = (
+                "Must pass in a non-empty structured output schema. Received: "
+                f"{schema_}"
+            )
+            raise ValueError(err_msg)
+        # Avoid mutating a caller-provided dict when merging extra kwargs.
+        structured_output_kwargs = dict(structured_output_kwargs or {})
+        for k in set(kwargs).difference(get_pydantic_field_names(self.__class__)):
+            structured_output_kwargs[k] = kwargs.pop(k)
+        super().__init__(
+            messages=messages,
+            schema_=schema_,
+            structured_output_kwargs=structured_output_kwargs,
+            template_format=template_format,
+            **kwargs,
+        )
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        For example, if the class is `coco.llms.openai.OpenAI`, then the namespace
+        is `["coco", "llms", "openai"]`
+
+        Returns:
+            The namespace of the Coco object.
+        """
+        return cls.__module__.split(".")
+
+    @classmethod
+    def from_messages_and_schema(
+        cls,
+        messages: Sequence[MessageLikeRepresentation],
+        schema: dict[str, Any] | type,
+        **kwargs: Any,
+    ) -> ChatPromptTemplate:
+        """Create a chat prompt template from a variety of message formats.
+
+        Examples:
+            Instantiation from a list of message templates:
+
+            ```python
+            from coco_core.prompts import StructuredPrompt
+
+
+            class OutputSchema(BaseModel):
+                name: str
+                value: int
+
+
+            template = StructuredPrompt(
+                [
+                    ("human", "Hello, how are you?"),
+                    ("ai", "I'm doing well, thanks!"),
+                    ("human", "That's good to hear."),
+                ],
+                OutputSchema,
+            )
+            ```
+
+        Args:
+            messages: Sequence of message representations.
+
+                A message can be represented using the following formats:
+
+                1. `BaseMessagePromptTemplate`
+                2. `BaseMessage`
+                3. 2-tuple of `(message type, template)`; e.g.,
+                    `("human", "{user_input}")`
+                4. 2-tuple of `(message class, template)`
+                5. A string which is shorthand for `("human", template)`; e.g.,
+                    `"{user_input}"`
+            schema: A dictionary representation of function call, or a Pydantic model.
+            **kwargs: Any additional kwargs to pass through to
+                `ChatModel.with_structured_output(schema, **kwargs)`.
+
+        Returns:
+            A structured prompt template
+        """
+        return cls(messages, schema, **kwargs)
+
+    @overload
+    def __or__(
+        self, other: Mapping[str, Any]
+    ) -> RunnableSerializable[dict[str, Any], dict[str, Any]]: ...
+
+    @overload
+    def __or__(
+        self,
+        other: Callable[[PromptValue], Runnable[PromptValue, Other]]
+        | Callable[[PromptValue], Awaitable[Runnable[PromptValue, Other]]],
+    ) -> RunnableSerializable[dict[str, Any], Other]: ...
+
+    @overload
+    def __or__(
+        self,
+        other: Runnable[PromptValue, Other]
+        | Callable[[Iterator[PromptValue]], Iterator[Other]]
+        | Callable[[AsyncIterator[PromptValue]], AsyncIterator[Other]]
+        | Callable[[PromptValue], Other],
+    ) -> RunnableSerializable[dict[str, Any], Other]: ...
+
+    @override
+    def __or__(
+        self,
+        other: Runnable[PromptValue, Other]
+        | Callable[[Iterator[PromptValue]], Iterator[Other]]
+        | Callable[[AsyncIterator[PromptValue]], AsyncIterator[Other]]
+        | Callable[[PromptValue], Other]
+        | Mapping[str, Runnable[PromptValue, Any] | Callable[[PromptValue], Any] | Any],
+    ) -> RunnableSerializable[dict[str, Any], Any]:
+        return self.pipe(other)
+
+    def pipe(
+        self,
+        *others: Runnable[Any, Other]
+        | Callable[[Iterator[Any]], Iterator[Other]]
+        | Callable[[AsyncIterator[Any]], AsyncIterator[Other]]
+        | Callable[[Any], Other]
+        | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other] | Any],
+        name: str | None = None,
+    ) -> RunnableSerializable[dict[str, Any], Other]:
+        """Pipe the structured prompt to a language model.
+
+        Args:
+            others: The language model to pipe the structured prompt to.
+            name: The name of the pipeline.
+
+        Returns:
+            A `RunnableSequence` object.
+
+        Raises:
+            NotImplementedError: If the first element of `others` is not a language
+                model.
+        """
+        if (others and isinstance(others[0], BaseLanguageModel)) or hasattr(
+            others[0], "with_structured_output"
+        ):
+            return RunnableSequence(
+                self,
+                others[0].with_structured_output(
+                    self.schema_, **self.structured_output_kwargs
+                ),
+                *others[1:],
+                name=name,
+            )
+        msg = "Structured prompts need to be piped to a language model."
+        raise NotImplementedError(msg)