coco-framework 0.1.0__py3-none-any.whl

coco_core/outputs/chat_result.py

@@ -0,0 +1,39 @@
+"""Chat result schema."""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+from coco_core.outputs.chat_generation import ChatGeneration
+
+
+class ChatResult(BaseModel):
+    """Use to represent the result of a chat model call with a single prompt.
+
+    This container is used internally by some implementations of chat model, it will
+    eventually be mapped to a more general `LLMResult` object, and  then projected into
+    an `AIMessage` object.
+
+    Coco users working with chat models will usually access information via
+    `AIMessage` (returned from runnable interfaces) or `LLMResult` (available via
+    callbacks). Please refer the `AIMessage` and `LLMResult` schema documentation for
+    more information.
+    """
+
+    generations: list[ChatGeneration]
+    """List of the chat generations.
+
+    Generations is a list to allow for multiple candidate generations for a single
+    input prompt.
+    """
+
+    llm_output: dict[str, Any] | None = None
+    """For arbitrary model provider-specific output.
+
+    This dictionary is a free-form dictionary that can contain any information that the
+    provider wants to return. It is not standardized and keys may vary by provider and
+    over time.
+
+    Users should generally avoid relying on this field and instead rely on accessing
+    relevant information from standardized fields present in `AIMessage`.
+    """

coco_core/outputs/generation.py

@@ -0,0 +1,80 @@
+"""Generation output schema."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from coco_core.load import Serializable
+from coco_core.utils._merge import merge_dicts
+
+
+class Generation(Serializable):
+    """A single text generation output.
+
+    Generation represents the response from an "old-fashioned" LLM (string-in,
+    string-out) that generates regular text (not chat messages).
+
+    This model is used internally by chat model and will eventually be mapped to a more
+    general `LLMResult` object, and then projected into an `AIMessage` object.
+
+    Coco users working with chat models will usually access information via
+    `AIMessage` (returned from runnable interfaces) or `LLMResult` (available via
+    callbacks). Please refer to `AIMessage` and `LLMResult` for more information.
+    """
+
+    text: str
+    """Generated text output."""
+
+    generation_info: dict[str, Any] | None = None
+    """Raw response from the provider.
+
+    May include things like the reason for finishing or token log probabilities.
+    """
+
+    type: Literal["Generation"] = "Generation"
+    """Type is used exclusively for serialization purposes.
+
+    Set to `'Generation'` for this class.
+    """
+
+    @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", "schema", "output"]`
+        """
+        return ["coco", "schema", "output"]
+
+
+class GenerationChunk(Generation):
+    """`GenerationChunk`, which can be concatenated with other `Generation` chunks."""
+
+    def __add__(self, other: GenerationChunk) -> GenerationChunk:
+        """Concatenate two `GenerationChunk` objects.
+
+        Args:
+            other: Another `GenerationChunk` to concatenate with.
+
+        Raises:
+            TypeError: If other is not a `GenerationChunk`.
+
+        Returns:
+            A new `GenerationChunk` concatenated from self and other.
+        """
+        if isinstance(other, GenerationChunk):
+            generation_info = merge_dicts(
+                self.generation_info or {},
+                other.generation_info or {},
+            )
+            return GenerationChunk(
+                text=self.text + other.text,
+                generation_info=generation_info or None,
+            )
+        msg = f"unsupported operand type(s) for +: '{type(self)}' and '{type(other)}'"  # type: ignore[unreachable]
+        raise TypeError(msg)

coco_core/outputs/llm_result.py

@@ -0,0 +1,112 @@
+"""`LLMResult` class."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+from coco_core.outputs.chat_generation import ChatGeneration, ChatGenerationChunk
+from coco_core.outputs.generation import Generation, GenerationChunk
+from coco_core.outputs.run_info import RunInfo
+
+
+class LLMResult(BaseModel):
+    """A container for results of an LLM call.
+
+    Both chat models and LLMs generate an `LLMResult` object. This object contains the
+    generated outputs and any additional information that the model provider wants to
+    return.
+    """
+
+    generations: list[
+        list[Generation | ChatGeneration | GenerationChunk | ChatGenerationChunk]
+    ]
+    """Generated outputs.
+
+    The first dimension of the list represents completions for different input prompts.
+
+    The second dimension of the list represents different candidate generations for a
+    given prompt.
+
+    - When returned from **an LLM**, the type is `list[list[Generation]]`.
+    - When returned from a **chat model**, the type is `list[list[ChatGeneration]]`.
+
+    `ChatGeneration` is a subclass of `Generation` that has a field for a structured
+    chat message.
+    """
+
+    llm_output: dict[str, Any] | None = None
+    """For arbitrary model provider-specific output.
+
+    This dictionary is a free-form dictionary that can contain any information that the
+    provider wants to return. It is not standardized and keys may vary by provider and
+    over time.
+
+    Users should generally avoid relying on this field and instead rely on accessing
+    relevant information from standardized fields present in AIMessage.
+    """
+
+    run: list[RunInfo] | None = None
+    """List of metadata info for model call for each input.
+
+    See `coco_core.outputs.run_info.RunInfo` for details.
+    """
+
+    type: Literal["LLMResult"] = "LLMResult"
+    """Type is used exclusively for serialization purposes."""
+
+    def flatten(self) -> list[LLMResult]:
+        """Flatten generations into a single list.
+
+        Unpack `list[list[Generation]] -> list[LLMResult]` where each returned
+        `LLMResult` contains only a single `Generation`. If token usage information is
+        available, it is kept only for the `LLMResult` corresponding to the top-choice
+        `Generation`, to avoid over-counting of token usage downstream.
+
+        Returns:
+            List of `LLMResult` objects where each returned `LLMResult` contains a
+                single `Generation`.
+        """
+        llm_results = []
+        for i, gen_list in enumerate(self.generations):
+            # Avoid double counting tokens in OpenAICallback
+            if i == 0:
+                llm_results.append(
+                    LLMResult(
+                        generations=[gen_list],
+                        llm_output=self.llm_output,
+                    )
+                )
+            else:
+                if self.llm_output is not None:
+                    llm_output = deepcopy(self.llm_output)
+                    llm_output["token_usage"] = {}
+                else:
+                    llm_output = None
+                llm_results.append(
+                    LLMResult(
+                        generations=[gen_list],
+                        llm_output=llm_output,
+                    )
+                )
+        return llm_results
+
+    def __eq__(self, other: object) -> bool:
+        """Check for `LLMResult` equality by ignoring any metadata related to runs.
+
+        Args:
+            other: Another `LLMResult` object to compare against.
+
+        Returns:
+            `True` if the generations and `llm_output` are equal, `False` otherwise.
+        """
+        if not isinstance(other, LLMResult):
+            return NotImplemented
+        return (
+            self.generations == other.generations
+            and self.llm_output == other.llm_output
+        )
+
+    __hash__ = None  # type: ignore[assignment]

coco_core/outputs/run_info.py

@@ -0,0 +1,22 @@
+"""`RunInfo` class."""
+
+from __future__ import annotations
+
+from uuid import UUID
+
+from pydantic import BaseModel
+
+
+class RunInfo(BaseModel):
+    """Class that contains metadata for a single execution of a chain or model.
+
+    Defined for backwards compatibility with older versions of `coco_core`.
+
+    !!! warning "This model will likely be deprecated in the future."
+
+    Users can acquire the `run_id` information from callbacks or via `run_id`
+    information present in the `astream_event` API (depending on the use case).
+    """
+
+    run_id: UUID
+    """A unique identifier for the model or chain run."""

coco_core/prompt_values.py

@@ -0,0 +1,161 @@
+"""**Prompt values** for language model prompts.
+
+Prompt values are used to represent different pieces of prompts. They can be used to
+represent text, images, or chat message pieces.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Any, Literal, cast
+
+from typing_extensions import TypedDict
+
+from coco_core.load.serializable import Serializable
+from coco_core.messages import (
+    AnyMessage,
+    BaseMessage,
+    HumanMessage,
+    get_buffer_string,
+)
+
+
+class PromptValue(Serializable, ABC):
+    """Base abstract class for inputs to any language model.
+
+    `PromptValues` can be converted to both LLM (pure text-generation) inputs and
+    chat model inputs.
+    """
+
+    @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", "schema", "prompt"]`
+        """
+        return ["coco", "schema", "prompt"]
+
+    @abstractmethod
+    def to_string(self) -> str:
+        """Return prompt value as string."""
+
+    @abstractmethod
+    def to_messages(self) -> list[BaseMessage]:
+        """Return prompt as a list of messages."""
+
+
+class StringPromptValue(PromptValue):
+    """String prompt value."""
+
+    text: str
+    """Prompt text."""
+
+    type: Literal["StringPromptValue"] = "StringPromptValue"
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "prompts", "base"]`
+        """
+        return ["coco", "prompts", "base"]
+
+    def to_string(self) -> str:
+        """Return prompt as string."""
+        return self.text
+
+    def to_messages(self) -> list[BaseMessage]:
+        """Return prompt as messages."""
+        return [HumanMessage(content=self.text)]
+
+
+class ChatPromptValue(PromptValue):
+    """Chat prompt value.
+
+    A type of a prompt value that is built from messages.
+    """
+
+    messages: Sequence[BaseMessage]
+    """List of messages."""
+
+    def to_string(self) -> str:
+        """Return prompt as string."""
+        return get_buffer_string(self.messages)
+
+    def to_messages(self) -> list[BaseMessage]:
+        """Return prompt as a list of messages."""
+        return list(self.messages)
+
+    @classmethod
+    def get_lc_namespace(cls) -> list[str]:
+        """Get the namespace of the Coco object.
+
+        Returns:
+            `["coco", "prompts", "chat"]`
+        """
+        return ["coco", "prompts", "chat"]
+
+
+class ImageURL(TypedDict, total=False):
+    """Image URL for multimodal model inputs (OpenAI format).
+
+    Represents the inner `image_url` object in OpenAI's Chat Completion API format. This
+    is used by `ImagePromptTemplate` and `ChatPromptTemplate`.
+
+    See Also:
+        `ImageContentBlock`: Coco's provider-agnostic image format used in message
+        content blocks. Use `ImageContentBlock` when working with the standardized
+        message format across different providers.
+
+    Note:
+        The `detail` field values are not validated locally. Invalid values
+        will be rejected by the downstream API, allowing new valid values to
+        be used without requiring a Coco update.
+    """
+
+    detail: Literal["auto", "low", "high"]
+    """Specifies the detail level of the image.
+
+    Defaults to ``'auto'`` if not specified. Higher detail levels consume
+    more tokens but provide better image understanding.
+    """
+
+    url: str
+    """URL of the image or base64-encoded image data."""
+
+
+class ImagePromptValue(PromptValue):
+    """Image prompt value."""
+
+    image_url: ImageURL
+    """Image URL."""
+
+    type: Literal["ImagePromptValue"] = "ImagePromptValue"
+
+    def to_string(self) -> str:
+        """Return prompt (image URL) as string."""
+        return self.image_url.get("url", "")
+
+    def to_messages(self) -> list[BaseMessage]:
+        """Return prompt (image URL) as messages."""
+        return [HumanMessage(content=[cast("dict[str, Any]", self.image_url)])]
+
+
+class ChatPromptValueConcrete(ChatPromptValue):
+    """Chat prompt value which explicitly lists out the message types it accepts.
+
+    For use in external schemas.
+    """
+
+    messages: Sequence[AnyMessage]
+    """Sequence of messages."""
+
+    type: Literal["ChatPromptValueConcrete"] = "ChatPromptValueConcrete"

coco_core/prompts/__init__.py

@@ -0,0 +1,101 @@
+"""A prompt is the input to the model.
+
+Prompt is often constructed from multiple components and prompt values. Prompt classes
+and functions make constructing and working with prompts easy.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core.prompts.base import (
+        BasePromptTemplate,
+        aformat_document,
+        format_document,
+    )
+    from coco_core.prompts.chat import (
+        AIMessagePromptTemplate,
+        BaseChatPromptTemplate,
+        ChatMessagePromptTemplate,
+        ChatPromptTemplate,
+        HumanMessagePromptTemplate,
+        MessagesPlaceholder,
+        SystemMessagePromptTemplate,
+    )
+    from coco_core.prompts.dict import DictPromptTemplate
+    from coco_core.prompts.few_shot import (
+        FewShotChatMessagePromptTemplate,
+        FewShotPromptTemplate,
+    )
+    from coco_core.prompts.few_shot_with_templates import (
+        FewShotPromptWithTemplates,
+    )
+    from coco_core.prompts.loading import load_prompt
+    from coco_core.prompts.prompt import PromptTemplate
+    from coco_core.prompts.string import (
+        StringPromptTemplate,
+        check_valid_template,
+        get_template_variables,
+        jinja2_formatter,
+        validate_jinja2,
+    )
+
+__all__ = (
+    "AIMessagePromptTemplate",
+    "BaseChatPromptTemplate",
+    "BasePromptTemplate",
+    "ChatMessagePromptTemplate",
+    "ChatPromptTemplate",
+    "DictPromptTemplate",
+    "FewShotChatMessagePromptTemplate",
+    "FewShotPromptTemplate",
+    "FewShotPromptWithTemplates",
+    "HumanMessagePromptTemplate",
+    "MessagesPlaceholder",
+    "PromptTemplate",
+    "StringPromptTemplate",
+    "SystemMessagePromptTemplate",
+    "aformat_document",
+    "check_valid_template",
+    "format_document",
+    "get_template_variables",
+    "jinja2_formatter",
+    "load_prompt",
+    "validate_jinja2",
+)
+
+_dynamic_imports = {
+    "BasePromptTemplate": "base",
+    "format_document": "base",
+    "aformat_document": "base",
+    "AIMessagePromptTemplate": "chat",
+    "BaseChatPromptTemplate": "chat",
+    "ChatMessagePromptTemplate": "chat",
+    "ChatPromptTemplate": "chat",
+    "DictPromptTemplate": "dict",
+    "HumanMessagePromptTemplate": "chat",
+    "MessagesPlaceholder": "chat",
+    "SystemMessagePromptTemplate": "chat",
+    "FewShotChatMessagePromptTemplate": "few_shot",
+    "FewShotPromptTemplate": "few_shot",
+    "FewShotPromptWithTemplates": "few_shot_with_templates",
+    "load_prompt": "loading",
+    "PromptTemplate": "prompt",
+    "StringPromptTemplate": "string",
+    "check_valid_template": "string",
+    "get_template_variables": "string",
+    "jinja2_formatter": "string",
+    "validate_jinja2": "string",
+}
+
+
+def __getattr__(attr_name: str) -> object:
+    module_name = _dynamic_imports.get(attr_name)
+    result = import_attr(attr_name, module_name, __spec__.parent)
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    return list(__all__)