coco-framework 0.1.0__py3-none-any.whl

coco_core/messages/base.py

@@ -0,0 +1,518 @@
+"""Base message."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast, overload
+
+from pydantic import ConfigDict, Field
+
+from coco_core._api.deprecation import warn_deprecated
+from coco_core.load.serializable import Serializable
+from coco_core.messages import content as types
+from coco_core.utils import get_bolded_text
+from coco_core.utils._merge import merge_dicts, merge_lists
+from coco_core.utils.interactive_env import is_interactive_env
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from typing_extensions import Self
+
+    from coco_core.prompts.chat import ChatPromptTemplate
+
+
+def _extract_reasoning_from_additional_kwargs(
+    message: BaseMessage,
+) -> types.ReasoningContentBlock | None:
+    """Extract `reasoning_content` from `additional_kwargs`.
+
+    Handles reasoning content stored in various formats:
+    - `additional_kwargs["reasoning_content"]` (string) - Ollama, DeepSeek, XAI, Groq
+
+    Args:
+        message: The message to extract reasoning from.
+
+    Returns:
+        A `ReasoningContentBlock` if reasoning content is found, None otherwise.
+    """
+    additional_kwargs = getattr(message, "additional_kwargs", {})
+
+    reasoning_content = additional_kwargs.get("reasoning_content")
+    if reasoning_content is not None and isinstance(reasoning_content, str):
+        return {"type": "reasoning", "reasoning": reasoning_content}
+
+    return None
+
+
+class TextAccessor(str):
+    """String-like object that supports both property and method access patterns.
+
+    Exists to maintain backward compatibility while transitioning from method-based to
+    property-based text access in message objects. In Coco <v1.0, message text was
+    accessed via `.text()` method calls. In v1.0=<, the preferred pattern is property
+    access via `.text`.
+
+    Rather than breaking existing code immediately, `TextAccessor` allows both
+    patterns:
+    - Modern property access: `message.text` (returns string directly)
+    - Legacy method access: `message.text()` (callable, emits deprecation warning)
+
+    """
+
+    __slots__ = ()
+
+    def __new__(cls, value: str) -> Self:
+        """Create new TextAccessor instance."""
+        return str.__new__(cls, value)
+
+    def __call__(self) -> str:
+        """Enable method-style text access for backward compatibility.
+
+        This method exists solely to support legacy code that calls `.text()`
+        as a method. New code should use property access (`.text`) instead.
+
+        !!! deprecated
+            As of `trustcoco-core` 1.0.0, calling `.text()` as a method is deprecated.
+            Use `.text` as a property instead. This method will be removed in 2.0.0.
+
+        Returns:
+            The string content, identical to property access.
+
+        """
+        warn_deprecated(
+            since="1.0.0",
+            message=(
+                "Calling .text() as a method is deprecated. "
+                "Use .text as a property instead (e.g., message.text)."
+            ),
+            removal="2.0.0",
+        )
+        return str(self)
+
+
+class BaseMessage(Serializable):
+    """Base abstract message class.
+
+    Messages are the inputs and outputs of a chat model.
+
+    Examples include [`HumanMessage`][coco.messages.HumanMessage],
+    [`AIMessage`][coco.messages.AIMessage], and
+    [`SystemMessage`][coco.messages.SystemMessage].
+    """
+
+    content: str | list[str | dict[Any, Any]]
+    """The contents of the message."""
+
+    additional_kwargs: dict[Any, Any] = Field(default_factory=dict)
+    """Reserved for additional payload data associated with the message.
+
+    For example, for a message from an AI, this could include tool calls as
+    encoded by the model provider.
+
+    """
+
+    response_metadata: dict[Any, Any] = Field(default_factory=dict)
+    """Examples: response headers, logprobs, token counts, model name."""
+
+    type: str
+    """The type of the message. Must be a string that is unique to the message type.
+
+    The purpose of this field is to allow for easy identification of the message type
+    when deserializing messages.
+
+    """
+
+    name: str | None = None
+    """An optional name for the message.
+
+    This can be used to provide a human-readable name for the message.
+
+    Usage of this field is optional, and whether it's used or not is up to the
+    model implementation.
+
+    """
+
+    id: str | None = Field(default=None, coerce_numbers_to_str=True)
+    """An optional unique identifier for the message.
+
+    This should ideally be provided by the provider/model which created the message.
+
+    """
+
+    model_config = ConfigDict(
+        extra="allow",
+    )
+
+    @overload
+    def __init__(
+        self,
+        content: str | list[str | dict[Any, Any]],
+        **kwargs: Any,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        content: str | list[str | dict[Any, Any]] | None = None,
+        content_blocks: list[types.ContentBlock] | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        content: str | list[str | dict[Any, Any]] | None = None,
+        content_blocks: list[types.ContentBlock] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a `BaseMessage`.
+
+        Specify `content` as positional arg or `content_blocks` for typing.
+
+        Args:
+            content: The contents of the message.
+            content_blocks: Typed standard content.
+            **kwargs: Additional arguments to pass to the parent class.
+        """
+        if content_blocks is not None:
+            super().__init__(content=content_blocks, **kwargs)
+        else:
+            super().__init__(content=content, **kwargs)
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """`BaseMessage` is serializable.
+
+        Returns:
+            True
+        """
+        return True
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "schema", "messages"]`
+        """
+        return ["coco", "schema", "messages"]
+
+    @property
+    def content_blocks(self) -> list[types.ContentBlock]:
+        r"""Load content blocks from the message content.
+
+        !!! version-added "Added in `trustcoco-core` 1.0.0"
+
+        """
+        # Needed here to avoid circular import, as these classes import BaseMessages
+        from coco_core.messages.block_translators.anthropic import (  # noqa: PLC0415
+            _convert_to_v1_from_anthropic_input,
+        )
+        from coco_core.messages.block_translators.bedrock_converse import (  # noqa: PLC0415
+            _convert_to_v1_from_converse_input,
+        )
+        from coco_core.messages.block_translators.google_genai import (  # noqa: PLC0415
+            _convert_to_v1_from_genai_input,
+        )
+        from coco_core.messages.block_translators.coco_v0 import (  # noqa: PLC0415
+            _convert_v0_multimodal_input_to_v1,
+        )
+        from coco_core.messages.block_translators.openai import (  # noqa: PLC0415
+            _convert_to_v1_from_chat_completions_input,
+        )
+
+        blocks: list[types.ContentBlock] = []
+        content = (
+            # Transpose string content to list, otherwise assumed to be list
+            [self.content]
+            if isinstance(self.content, str) and self.content
+            else self.content
+        )
+        for item in content:
+            if isinstance(item, str):
+                # Plain string content is treated as a text block
+                blocks.append({"type": "text", "text": item})
+            elif isinstance(item, dict):
+                item_type = item.get("type")
+                if item_type not in types.KNOWN_BLOCK_TYPES:
+                    # Handle all provider-specific or None type blocks as non-standard -
+                    # we'll come back to these later
+                    blocks.append({"type": "non_standard", "value": item})
+                else:
+                    # Guard against v0 blocks that share the same `type` keys
+                    if "source_type" in item:
+                        blocks.append({"type": "non_standard", "value": item})
+                        continue
+
+                    # This can't be a v0 block (since they require `source_type`),
+                    # so it's a known v1 block type
+                    blocks.append(cast("types.ContentBlock", item))
+
+        # Subsequent passes: attempt to unpack non-standard blocks.
+        # This is the last stop - if we can't parse it here, it is left as non-standard
+        for parsing_step in [
+            _convert_v0_multimodal_input_to_v1,
+            _convert_to_v1_from_chat_completions_input,
+            _convert_to_v1_from_anthropic_input,
+            _convert_to_v1_from_genai_input,
+            _convert_to_v1_from_converse_input,
+        ]:
+            blocks = parsing_step(blocks)
+        return blocks
+
+    @property
+    def text(self) -> TextAccessor:
+        """Get the text content of the message as a string.
+
+        Can be used as both property (`message.text`) and method (`message.text()`).
+
+        Handles both string and list content types (e.g. for content blocks). Only
+        extracts blocks with `type: 'text'`; other block types are ignored.
+
+        !!! deprecated
+            As of `trustcoco-core` 1.0.0, calling `.text()` as a method is deprecated.
+            Use `.text` as a property instead. This method will be removed in 2.0.0.
+
+        Returns:
+            The text content of the message.
+
+        """
+        if isinstance(self.content, str):
+            text_value = self.content
+        else:
+            # Must be a list
+            blocks = [
+                block
+                for block in self.content
+                if isinstance(block, str)
+                or (block.get("type") == "text" and isinstance(block.get("text"), str))
+            ]
+            text_value = "".join(
+                block if isinstance(block, str) else block["text"] for block in blocks
+            )
+        return TextAccessor(text_value)
+
+    def __add__(self, other: Any) -> ChatPromptTemplate:
+        """Concatenate this message with another message.
+
+        Args:
+            other: Another message to concatenate with this one.
+
+        Returns:
+            A ChatPromptTemplate containing both messages.
+        """
+        # Import locally to prevent circular imports.
+        from coco_core.prompts.chat import ChatPromptTemplate  # noqa: PLC0415
+
+        prompt = ChatPromptTemplate(messages=[self])
+        return prompt.__add__(other)
+
+    def pretty_repr(
+        self,
+        html: bool = False,  # noqa: FBT001,FBT002
+    ) -> str:
+        """Get a pretty representation of the message.
+
+        Args:
+            html: Whether to format the message as HTML. If `True`, the message will be
+                formatted with HTML tags.
+
+        Returns:
+            A pretty representation of the message.
+
+        Example:
+            ```python
+            from coco_core.messages import HumanMessage
+
+            msg = HumanMessage(content="What is the capital of France?")
+            print(msg.pretty_repr())
+            ```
+
+            Results in:
+
+            ```txt
+            ================================ Human Message =================================
+
+            What is the capital of France?
+            ```
+        """  # noqa: E501
+        title = get_msg_title_repr(self.type.title() + " Message", bold=html)
+        # TODO: handle non-string content.
+        if self.name is not None:
+            title += f"\nName: {self.name}"
+        return f"{title}\n\n{self.content}"
+
+    def pretty_print(self) -> None:
+        """Print a pretty representation of the message.
+
+        Example:
+            ```python
+            from coco_core.messages import AIMessage
+
+            msg = AIMessage(content="The capital of France is Paris.")
+            msg.pretty_print()
+            ```
+
+            Results in:
+
+            ```txt
+            ================================== Ai Message ==================================
+
+            The capital of France is Paris.
+            ```
+        """  # noqa: E501
+        print(self.pretty_repr(html=is_interactive_env()))  # noqa: T201
+
+
+def merge_content(
+    first_content: str | list[str | dict[Any, Any]],
+    *contents: str | list[str | dict[Any, Any]],
+) -> str | list[str | dict[Any, Any]]:
+    """Merge multiple message contents.
+
+    Args:
+        first_content: The first `content`. Can be a string or a list.
+        contents: The other `content`s. Can be a string or a list.
+
+    Returns:
+        The merged content.
+
+    """
+    merged: str | list[str | dict[Any, Any]]
+    merged = "" if first_content is None else first_content
+
+    for content in contents:
+        # If current is a string
+        if isinstance(merged, str):
+            # If the next chunk is also a string, then merge them naively
+            if isinstance(content, str):
+                merged += content
+            # If the next chunk is a list, add the current to the start of the list
+            else:
+                merged = [merged, *content]
+        elif isinstance(content, list):
+            # If both are lists
+            merged = merge_lists(merged, content)  # type: ignore[assignment]
+        # If the first content is a list, and the second content is a string
+        # If the last element of the first content is a string
+        # Add the second content to the last element
+        elif merged and isinstance(merged[-1], str):
+            merged[-1] += content
+        # If second content is an empty string, treat as a no-op
+        elif content == "":
+            pass
+        # Otherwise, add the second content as a new element of the list
+        elif merged:
+            merged.append(content)
+    return merged
+
+
+class BaseMessageChunk(BaseMessage):
+    """Message chunk, which can be concatenated with other Message chunks."""
+
+    def __add__(self, other: Any) -> BaseMessageChunk:  # type: ignore[override]
+        """Message chunks support concatenation with other message chunks.
+
+        This functionality is useful to combine message chunks yielded from
+        a streaming model into a complete message.
+
+        Args:
+            other: Another message chunk to concatenate with this one.
+
+        Returns:
+            A new message chunk that is the concatenation of this message chunk
+            and the other message chunk.
+
+        Raises:
+            TypeError: If the other object is not a message chunk.
+
+        Example:
+            ```txt
+              AIMessageChunk(content="Hello", ...)
+            + AIMessageChunk(content=" World", ...)
+            = AIMessageChunk(content="Hello World", ...)
+            ```
+        """
+        if isinstance(other, BaseMessageChunk):
+            # If both are (subclasses of) BaseMessageChunk,
+            # concat into a single BaseMessageChunk
+
+            return self.__class__(
+                id=self.id,
+                type=self.type,
+                content=merge_content(self.content, other.content),
+                additional_kwargs=merge_dicts(
+                    self.additional_kwargs, other.additional_kwargs
+                ),
+                response_metadata=merge_dicts(
+                    self.response_metadata, other.response_metadata
+                ),
+            )
+        if isinstance(other, list) and all(
+            isinstance(o, BaseMessageChunk) for o in other
+        ):
+            content = merge_content(self.content, *(o.content for o in other))
+            additional_kwargs = merge_dicts(
+                self.additional_kwargs, *(o.additional_kwargs for o in other)
+            )
+            response_metadata = merge_dicts(
+                self.response_metadata, *(o.response_metadata for o in other)
+            )
+            return self.__class__(  # type: ignore[call-arg]
+                id=self.id,
+                content=content,
+                additional_kwargs=additional_kwargs,
+                response_metadata=response_metadata,
+            )
+        msg = (
+            'unsupported operand type(s) for +: "'
+            f"{self.__class__.__name__}"
+            f'" and "{other.__class__.__name__}"'
+        )
+        raise TypeError(msg)
+
+
+def message_to_dict(message: BaseMessage) -> dict[str, Any]:
+    """Convert a Message to a dictionary.
+
+    Args:
+        message: Message to convert.
+
+    Returns:
+        Message as a dict. The dict will have a `type` key with the message type
+        and a `data` key with the message data as a dict.
+
+    """
+    return {"type": message.type, "data": message.model_dump()}
+
+
+def messages_to_dict(messages: Sequence[BaseMessage]) -> list[dict[str, Any]]:
+    """Convert a sequence of Messages to a list of dictionaries.
+
+    Args:
+        messages: Sequence of messages (as `BaseMessage`s) to convert.
+
+    Returns:
+        List of messages as dicts.
+
+    """
+    return [message_to_dict(m) for m in messages]
+
+
+def get_msg_title_repr(title: str, *, bold: bool = False) -> str:
+    """Get a title representation for a message.
+
+    Args:
+        title: The title.
+        bold: Whether to bold the title.
+
+    Returns:
+        The title representation.
+
+    """
+    padded = " " + title + " "
+    sep_len = (80 - len(padded)) // 2
+    sep = "=" * sep_len
+    second_sep = sep + "=" if len(padded) % 2 else sep
+    if bold:
+        padded = get_bolded_text(padded)
+    return f"{sep}{padded}{second_sep}"

coco_core/messages/block_translators/__init__.py

@@ -0,0 +1,112 @@
+"""Derivations of standard content blocks from provider content.
+
+`AIMessage` will first attempt to use a provider-specific translator if
+`model_provider` is set in `response_metadata` on the message. Consequently, each
+provider translator must handle all possible content response types from the provider,
+including text.
+
+If no provider is set, or if the provider does not have a registered translator,
+`AIMessage` will fall back to best-effort parsing of the content into blocks using
+the implementation in `BaseMessage`.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from coco_core.messages import AIMessage, AIMessageChunk
+    from coco_core.messages import content as types
+
+# Provider to translator mapping
+PROVIDER_TRANSLATORS: dict[str, dict[str, Callable[..., list[types.ContentBlock]]]] = {}
+"""Map model provider names to translator functions.
+
+The dictionary maps provider names (e.g. `'openai'`, `'anthropic'`) to another
+dictionary with two keys:
+- `'translate_content'`: Function to translate `AIMessage` content.
+- `'translate_content_chunk'`: Function to translate `AIMessageChunk` content.
+
+When calling `content_blocks` on an `AIMessage` or `AIMessageChunk`, if
+`model_provider` is set in `response_metadata`, the corresponding translator
+functions will be used to parse the content into blocks. Otherwise, best-effort parsing
+in `BaseMessage` will be used.
+"""
+
+
+def register_translator(
+    provider: str,
+    translate_content: Callable[[AIMessage], list[types.ContentBlock]],
+    translate_content_chunk: Callable[[AIMessageChunk], list[types.ContentBlock]],
+) -> None:
+    """Register content translators for a provider in `PROVIDER_TRANSLATORS`.
+
+    Args:
+        provider: The model provider name (e.g. `'openai'`, `'anthropic'`).
+        translate_content: Function to translate `AIMessage` content.
+        translate_content_chunk: Function to translate `AIMessageChunk` content.
+    """
+    PROVIDER_TRANSLATORS[provider] = {
+        "translate_content": translate_content,
+        "translate_content_chunk": translate_content_chunk,
+    }
+
+
+def get_translator(
+    provider: str,
+) -> dict[str, Callable[..., list[types.ContentBlock]]] | None:
+    """Get the translator functions for a provider.
+
+    Args:
+        provider: The model provider name.
+
+    Returns:
+        Dictionary with `'translate_content'` and `'translate_content_chunk'`
+        functions, or None if no translator is registered for the provider. In such
+        case, best-effort parsing in `BaseMessage` will be used.
+    """
+    return PROVIDER_TRANSLATORS.get(provider)
+
+
+def _register_translators() -> None:
+    """Register all translators in trustcoco-core.
+
+    A unit test ensures all modules in `block_translators` are represented here.
+
+    For translators implemented outside trustcoco-core, they can be registered by
+    calling `register_translator` from within the integration package.
+    """
+    from coco_core.messages.block_translators.anthropic import (  # noqa: PLC0415
+        _register_anthropic_translator,
+    )
+    from coco_core.messages.block_translators.bedrock import (  # noqa: PLC0415
+        _register_bedrock_translator,
+    )
+    from coco_core.messages.block_translators.bedrock_converse import (  # noqa: PLC0415
+        _register_bedrock_converse_translator,
+    )
+    from coco_core.messages.block_translators.google_genai import (  # noqa: PLC0415
+        _register_google_genai_translator,
+    )
+    from coco_core.messages.block_translators.google_vertexai import (  # noqa: PLC0415
+        _register_google_vertexai_translator,
+    )
+    from coco_core.messages.block_translators.groq import (  # noqa: PLC0415
+        _register_groq_translator,
+    )
+    from coco_core.messages.block_translators.openai import (  # noqa: PLC0415
+        _register_openai_translator,
+    )
+
+    _register_bedrock_translator()
+    _register_bedrock_converse_translator()
+    _register_anthropic_translator()
+    _register_google_genai_translator()
+    _register_google_vertexai_translator()
+    _register_groq_translator()
+    _register_openai_translator()
+
+
+_register_translators()