coco-framework 0.1.0__py3-none-any.whl

coco_core/callbacks/stdout.py

@@ -0,0 +1,123 @@
+"""Callback handler that prints to std out."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import override
+
+from coco_core.callbacks.base import BaseCallbackHandler
+from coco_core.utils import print_text
+
+if TYPE_CHECKING:
+    from coco_core.agents import AgentAction, AgentFinish
+
+
+class StdOutCallbackHandler(BaseCallbackHandler):
+    """Callback handler that prints to std out."""
+
+    def __init__(self, color: str | None = None) -> None:
+        """Initialize callback handler.
+
+        Args:
+            color: The color to use for the text.
+        """
+        self.color = color
+
+    @override
+    def on_chain_start(
+        self, serialized: dict[str, Any], inputs: dict[str, Any], **kwargs: Any
+    ) -> None:
+        """Print out that we are entering a chain.
+
+        Args:
+            serialized: The serialized chain.
+            inputs: The inputs to the chain.
+            **kwargs: Additional keyword arguments.
+        """
+        if "name" in kwargs:
+            name = kwargs["name"]
+        elif serialized:
+            name = serialized.get("name", serialized.get("id", ["<unknown>"])[-1])
+        else:
+            name = "<unknown>"
+        print(f"\n\n\033[1m> Entering new {name} chain...\033[0m")  # noqa: T201
+
+    @override
+    def on_chain_end(self, outputs: dict[str, Any], **kwargs: Any) -> None:
+        """Print out that we finished a chain.
+
+        Args:
+            outputs: The outputs of the chain.
+            **kwargs: Additional keyword arguments.
+        """
+        print("\n\033[1m> Finished chain.\033[0m")  # noqa: T201
+
+    @override
+    def on_agent_action(
+        self, action: AgentAction, color: str | None = None, **kwargs: Any
+    ) -> Any:
+        """Run on agent action.
+
+        Args:
+            action: The agent action.
+            color: The color to use for the text.
+            **kwargs: Additional keyword arguments.
+        """
+        print_text(action.log, color=color or self.color)
+
+    @override
+    def on_tool_end(
+        self,
+        output: Any,
+        color: str | None = None,
+        observation_prefix: str | None = None,
+        llm_prefix: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """If not the final action, print out observation.
+
+        Args:
+            output: The output to print.
+            color: The color to use for the text.
+            observation_prefix: The observation prefix.
+            llm_prefix: The LLM prefix.
+            **kwargs: Additional keyword arguments.
+        """
+        output = str(output)
+        if observation_prefix is not None:
+            print_text(f"\n{observation_prefix}")
+        print_text(output, color=color or self.color)
+        if llm_prefix is not None:
+            print_text(f"\n{llm_prefix}")
+
+    @override
+    def on_text(
+        self,
+        text: str,
+        color: str | None = None,
+        end: str = "",
+        **kwargs: Any,
+    ) -> None:
+        """Run when the agent ends.
+
+        Args:
+            text: The text to print.
+            color: The color to use for the text.
+            end: The end character to use.
+            **kwargs: Additional keyword arguments.
+        """
+        print_text(text, color=color or self.color, end=end)
+
+    @override
+    def on_agent_finish(
+        self, finish: AgentFinish, color: str | None = None, **kwargs: Any
+    ) -> None:
+        """Run on the agent end.
+
+        Args:
+            finish: The agent finish.
+            color: The color to use for the text.
+            **kwargs: Additional keyword arguments.
+        """
+        print_text(finish.log, color=color or self.color, end="\n")

coco_core/callbacks/streaming_stdout.py

@@ -0,0 +1,154 @@
+"""Callback Handler streams to stdout on new llm token."""
+
+from __future__ import annotations
+
+import sys
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import override
+
+from coco_core.callbacks.base import BaseCallbackHandler
+
+if TYPE_CHECKING:
+    from coco_core.agents import AgentAction, AgentFinish
+    from coco_core.messages import BaseMessage
+    from coco_core.outputs import LLMResult
+
+
+class StreamingStdOutCallbackHandler(BaseCallbackHandler):
+    """Callback handler for streaming.
+
+    !!! warning "Only works with LLMs that support streaming."
+    """
+
+    def on_llm_start(
+        self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any
+    ) -> None:
+        """Run when LLM starts running.
+
+        Args:
+            serialized: The serialized LLM.
+            prompts: The prompts to run.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[BaseMessage]],
+        **kwargs: Any,
+    ) -> None:
+        """Run when LLM starts running.
+
+        Args:
+            serialized: The serialized LLM.
+            messages: The messages to run.
+            **kwargs: Additional keyword arguments.
+        """
+
+    @override
+    def on_llm_new_token(
+        self, token: str | list[str | dict[str, Any]], **kwargs: Any
+    ) -> None:
+        """Run on new LLM token. Only available when streaming is enabled.
+
+        Args:
+            token: The new token, or a list of content blocks.
+            **kwargs: Additional keyword arguments.
+        """
+        sys.stdout.write(str(token))
+        sys.stdout.flush()
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Run when LLM ends running.
+
+        Args:
+            response: The response from the LLM.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_llm_error(self, error: BaseException, **kwargs: Any) -> None:
+        """Run when LLM errors.
+
+        Args:
+            error: The error that occurred.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_chain_start(
+        self, serialized: dict[str, Any], inputs: dict[str, Any], **kwargs: Any
+    ) -> None:
+        """Run when a chain starts running.
+
+        Args:
+            serialized: The serialized chain.
+            inputs: The inputs to the chain.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_chain_end(self, outputs: dict[str, Any], **kwargs: Any) -> None:
+        """Run when a chain ends running.
+
+        Args:
+            outputs: The outputs of the chain.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_chain_error(self, error: BaseException, **kwargs: Any) -> None:
+        """Run when chain errors.
+
+        Args:
+            error: The error that occurred.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_tool_start(
+        self, serialized: dict[str, Any], input_str: str, **kwargs: Any
+    ) -> None:
+        """Run when the tool starts running.
+
+        Args:
+            serialized: The serialized tool.
+            input_str: The input string.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
+        """Run on agent action.
+
+        Args:
+            action: The agent action.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_tool_end(self, output: Any, **kwargs: Any) -> None:
+        """Run when tool ends running.
+
+        Args:
+            output: The output of the tool.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_tool_error(self, error: BaseException, **kwargs: Any) -> None:
+        """Run when tool errors.
+
+        Args:
+            error: The error that occurred.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_text(self, text: str, **kwargs: Any) -> None:
+        """Run on an arbitrary text.
+
+        Args:
+            text: The text to print.
+            **kwargs: Additional keyword arguments.
+        """
+
+    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
+        """Run on the agent end.
+
+        Args:
+            finish: The agent finish.
+            **kwargs: Additional keyword arguments.
+        """

coco_core/callbacks/usage.py

@@ -0,0 +1,124 @@
+"""Callback Handler that tracks `AIMessage.usage_metadata`."""
+
+import threading
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any
+
+from typing_extensions import override
+
+from coco_core.callbacks import BaseCallbackHandler
+from coco_core.messages import AIMessage
+from coco_core.messages.ai import UsageMetadata, add_usage
+from coco_core.outputs import ChatGeneration, LLMResult
+from coco_core.tracers.context import register_configure_hook
+
+
+class UsageMetadataCallbackHandler(BaseCallbackHandler):
+    """Callback Handler that tracks `AIMessage.usage_metadata`.
+
+    Example:
+        ```python
+        from coco.chat_models import init_chat_model
+        from coco_core.callbacks import UsageMetadataCallbackHandler
+
+        llm_1 = init_chat_model(model="openai:gpt-5.5")
+        llm_2 = init_chat_model(model="anthropic:claude-haiku-4-5-20251001")
+
+        callback = UsageMetadataCallbackHandler()
+        result_1 = llm_1.invoke("Hello", config={"callbacks": [callback]})
+        result_2 = llm_2.invoke("Hello", config={"callbacks": [callback]})
+        callback.usage_metadata
+        ```
+
+    !!! version-added "Added in `trustcoco-core` 0.3.49"
+
+    """
+
+    def __init__(self) -> None:
+        """Initialize the `UsageMetadataCallbackHandler`."""
+        super().__init__()
+        self._lock = threading.Lock()
+        self.usage_metadata: dict[str, UsageMetadata] = {}
+
+    @override
+    def __repr__(self) -> str:
+        return str(self.usage_metadata)
+
+    @override
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Collect token usage."""
+        # Check for usage_metadata (trustcoco-core >= 0.2.2)
+        try:
+            generation = response.generations[0][0]
+        except IndexError:
+            generation = None
+
+        usage_metadata = None
+        model_name = None
+        if isinstance(generation, ChatGeneration):
+            try:
+                message = generation.message
+                if isinstance(message, AIMessage):
+                    usage_metadata = message.usage_metadata
+                    model_name = message.response_metadata.get("model_name")
+            except AttributeError:
+                pass
+
+        # update shared state behind lock
+        if usage_metadata and model_name:
+            with self._lock:
+                if model_name not in self.usage_metadata:
+                    self.usage_metadata[model_name] = usage_metadata
+                else:
+                    self.usage_metadata[model_name] = add_usage(
+                        self.usage_metadata[model_name], usage_metadata
+                    )
+
+
+@contextmanager
+def get_usage_metadata_callback(
+    name: str = "usage_metadata_callback",
+) -> Generator[UsageMetadataCallbackHandler, None, None]:
+    """Get usage metadata callback.
+
+    Get context manager for tracking usage metadata across chat model calls using
+    [`AIMessage.usage_metadata`][coco.messages.AIMessage.usage_metadata].
+
+    Args:
+        name: The name of the context variable.
+
+    Yields:
+        The usage metadata callback.
+
+    Example:
+        ```python
+        from coco.chat_models import init_chat_model
+        from coco_core.callbacks import get_usage_metadata_callback
+
+        llm_1 = init_chat_model(model="openai:gpt-5.5")
+        llm_2 = init_chat_model(model="anthropic:claude-haiku-4-5-20251001")
+
+        with get_usage_metadata_callback() as cb:
+            llm_1.invoke("Hello")
+            llm_2.invoke("Hello")
+            print(cb.usage_metadata)
+        ```
+
+    !!! version-added "Added in `trustcoco-core` 0.3.49"
+
+    """
+    usage_metadata_callback_var: ContextVar[UsageMetadataCallbackHandler | None] = (
+        ContextVar(name, default=None)
+    )
+    register_configure_hook(usage_metadata_callback_var, inheritable=True)
+    cb = UsageMetadataCallbackHandler()
+    token = usage_metadata_callback_var.set(cb)
+    try:
+        yield cb
+    finally:
+        # Always clear the context var, including when the with-block raises.
+        # Without finally, post-block model calls keep accumulating into cb
+        # (see #38989).
+        usage_metadata_callback_var.reset(token)

coco_core/chat_history.py

@@ -0,0 +1,267 @@
+"""Chat message history stores a history of the message interactions in a chat."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, Field
+
+from coco_core._api import deprecated
+from coco_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+    get_buffer_string,
+)
+from coco_core.runnables.config import run_in_executor
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+@deprecated(
+    since="1.6.4",
+    removal="2.0.0",
+    addendum=(
+        "See the short-term memory documentation for recommended alternatives: "
+        "https://docs.langchain.com/oss/python/langchain/short-term-memory"
+    ),
+)
+class BaseChatMessageHistory(ABC):
+    """Abstract base class for storing chat message history.
+
+    Implementations guidelines:
+
+    Implementations are expected to over-ride all or some of the following methods:
+
+    * `add_messages`: sync variant for bulk addition of messages
+    * `aadd_messages`: async variant for bulk addition of messages
+    * `messages`: sync variant for getting messages
+    * `aget_messages`: async variant for getting messages
+    * `clear`: sync variant for clearing messages
+    * `aclear`: async variant for clearing messages
+
+    `add_messages` contains a default implementation that calls `add_message`
+    for each message in the sequence. This is provided for backwards compatibility
+    with existing implementations which only had `add_message`.
+
+    Async variants all have default implementations that call the sync variants.
+    Implementers can choose to override the async implementations to provide
+    truly async implementations.
+
+    Usage guidelines:
+
+    When used for updating history, users should favor usage of `add_messages`
+    over `add_message` or other variants like `add_user_message` and `add_ai_message`
+    to avoid unnecessary round-trips to the underlying persistence layer.
+
+    Example:
+        ```python
+        import json
+        import os
+        from coco_core.messages import messages_from_dict, message_to_dict
+
+
+        class FileChatMessageHistory(BaseChatMessageHistory):
+            storage_path: str
+            session_id: str
+
+            @property
+            def messages(self) -> list[BaseMessage]:
+                try:
+                    with open(
+                        os.path.join(self.storage_path, self.session_id),
+                        "r",
+                        encoding="utf-8",
+                    ) as f:
+                        messages_data = json.load(f)
+                    return messages_from_dict(messages_data)
+                except FileNotFoundError:
+                    return []
+
+            def add_messages(self, messages: Sequence[BaseMessage]) -> None:
+                all_messages = list(self.messages)  # Existing messages
+                all_messages.extend(messages)  # Add new messages
+
+                serialized = [message_to_dict(message) for message in all_messages]
+                file_path = os.path.join(self.storage_path, self.session_id)
+                os.makedirs(os.path.dirname(file_path), exist_ok=True)
+                with open(file_path, "w", encoding="utf-8") as f:
+                    json.dump(serialized, f)
+
+            def clear(self) -> None:
+                file_path = os.path.join(self.storage_path, self.session_id)
+                os.makedirs(os.path.dirname(file_path), exist_ok=True)
+                with open(file_path, "w", encoding="utf-8") as f:
+                    json.dump([], f)
+        ```
+    """
+
+    messages: list[BaseMessage]
+    """A property or attribute that returns a list of messages.
+
+    In general, getting the messages may involve IO to the underlying persistence
+    layer, so this operation is expected to incur some latency.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize cooperatively to preserve multiple inheritance."""
+        super().__init__(*args, **kwargs)
+
+    async def aget_messages(self) -> list[BaseMessage]:
+        """Async version of getting messages.
+
+        Can over-ride this method to provide an efficient async implementation.
+
+        In general, fetching messages may involve IO to the underlying persistence
+        layer.
+
+        Returns:
+            The messages.
+        """
+        return await run_in_executor(None, lambda: self.messages)
+
+    def add_user_message(self, message: HumanMessage | str) -> None:
+        """Convenience method for adding a human message string to the store.
+
+        !!! note
+
+            This is a convenience method. Code should favor the bulk `add_messages`
+            interface instead to save on round-trips to the persistence layer.
+
+        This method may be deprecated in a future release.
+
+        Args:
+            message: The `HumanMessage` to add to the store.
+        """
+        if isinstance(message, HumanMessage):
+            self.add_message(message)
+        else:
+            self.add_message(HumanMessage(content=message))
+
+    def add_ai_message(self, message: AIMessage | str) -> None:
+        """Convenience method for adding an `AIMessage` string to the store.
+
+        !!! note
+
+            This is a convenience method. Code should favor the bulk `add_messages`
+            interface instead to save on round-trips to the persistence layer.
+
+        This method may be deprecated in a future release.
+
+        Args:
+            message: The `AIMessage` to add.
+        """
+        if isinstance(message, AIMessage):
+            self.add_message(message)
+        else:
+            self.add_message(AIMessage(content=message))
+
+    def add_message(self, message: BaseMessage) -> None:
+        """Add a Message object to the store.
+
+        Args:
+            message: A `BaseMessage` object to store.
+
+        Raises:
+            NotImplementedError: If the sub-class has not implemented an efficient
+                `add_messages` method.
+        """
+        if type(self).add_messages != BaseChatMessageHistory.add_messages:
+            # This means that the sub-class has implemented an efficient add_messages
+            # method, so we should use it.
+            self.add_messages([message])
+        else:
+            msg = (
+                "add_message is not implemented for this class. "
+                "Please implement add_message or add_messages."
+            )
+            raise NotImplementedError(msg)
+
+    def add_messages(self, messages: Sequence[BaseMessage]) -> None:
+        """Add a list of messages.
+
+        Implementations should over-ride this method to handle bulk addition of messages
+        in an efficient manner to avoid unnecessary round-trips to the underlying store.
+
+        Args:
+            messages: A sequence of `BaseMessage` objects to store.
+        """
+        for message in messages:
+            self.add_message(message)
+
+    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
+        """Async add a list of messages.
+
+        Args:
+            messages: A sequence of `BaseMessage` objects to store.
+        """
+        await run_in_executor(None, self.add_messages, messages)
+
+    @abstractmethod
+    def clear(self) -> None:
+        """Remove all messages from the store."""
+
+    async def aclear(self) -> None:
+        """Async remove all messages from the store."""
+        await run_in_executor(None, self.clear)
+
+    def __str__(self) -> str:
+        """Return a string representation of the chat history."""
+        return get_buffer_string(self.messages)
+
+
+@deprecated(
+    since="1.6.4",
+    removal="2.0.0",
+    addendum=(
+        "See the short-term memory documentation for recommended alternatives: "
+        "https://docs.langchain.com/oss/python/langchain/short-term-memory"
+    ),
+)
+class InMemoryChatMessageHistory(BaseChatMessageHistory, BaseModel):
+    """In memory implementation of chat message history.
+
+    Stores messages in a memory list.
+    """
+
+    messages: list[BaseMessage] = Field(default_factory=list)
+    """A list of messages stored in memory."""
+
+    async def aget_messages(self) -> list[BaseMessage]:
+        """Async version of getting messages.
+
+        Can over-ride this method to provide an efficient async implementation.
+
+        In general, fetching messages may involve IO to the underlying persistence
+        layer.
+
+        Returns:
+            List of messages.
+        """
+        return self.messages
+
+    def add_message(self, message: BaseMessage) -> None:
+        """Add a self-created message to the store.
+
+        Args:
+            message: The message to add.
+        """
+        self.messages.append(message)
+
+    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
+        """Async add messages to the store.
+
+        Args:
+            messages: The messages to add.
+        """
+        self.add_messages(messages)
+
+    def clear(self) -> None:
+        """Clear all messages from the store."""
+        self.messages = []
+
+    async def aclear(self) -> None:
+        """Async clear all messages from the store."""
+        self.clear()

coco_core/chat_loaders.py

@@ -0,0 +1,26 @@
+"""Chat loaders."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterator
+
+from coco_core.chat_sessions import ChatSession
+
+
+class BaseChatLoader(ABC):
+    """Base class for chat loaders."""
+
+    @abstractmethod
+    def lazy_load(self) -> Iterator[ChatSession]:
+        """Lazy load the chat sessions.
+
+        Returns:
+            An iterator of chat sessions.
+        """
+
+    def load(self) -> list[ChatSession]:
+        """Eagerly load the chat sessions into memory.
+
+        Returns:
+            A list of chat sessions.
+        """
+        return list(self.lazy_load())