coco-framework 0.1.0__py3-none-any.whl

coco_core/tracers/memory_stream.py

@@ -0,0 +1,200 @@
+"""Module implements a memory stream for communication between two co-routines.
+
+This module provides a way to communicate between two co-routines using a memory
+channel. The writer and reader can be in the same event loop or in different event
+loops. When they're in different event loops, they will also be in different threads.
+
+Useful in situations when there's a mix of synchronous and asynchronous used in the
+code.
+"""
+
+import asyncio
+import contextlib
+import weakref
+from asyncio import AbstractEventLoop, Queue
+from collections.abc import AsyncIterator
+from typing import Any, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+def _close_loop_quietly(loop: AbstractEventLoop) -> None:
+    """Close an event loop that was created internally, as a finalizer callback.
+
+    Errors are suppressed rather than logged. `weakref.finalize` defaults to
+    `atexit=True`, so this can run during interpreter shutdown, where `logging`
+    may already be torn down. Python routes finalizer exceptions to
+    `sys.unraisablehook` rather than propagating them, so suppressing here only
+    avoids dumping a traceback to stderr at a nondeterministic moment, which is
+    the same kind of noise this cleanup exists to remove. If `close()` does
+    fail, the loop's own `__del__` still emits
+    `ResourceWarning: unclosed event loop`, so the leak stays detectable.
+
+    Args:
+        loop: The event loop to close. A running loop is left alone.
+    """
+    if loop.is_running():
+        # `close()` would raise; leave the `ResourceWarning` as the signal.
+        return
+    with contextlib.suppress(Exception):
+        loop.close()
+
+
+def _get_or_create_loop(owner: object) -> AbstractEventLoop:
+    """Get the current event loop, creating one if this thread has none.
+
+    Args:
+        owner: Object whose lifetime bounds a loop created here.
+
+            A loop created by this function is closed once `owner` is
+            garbage collected.
+
+    Returns:
+        The event loop to back a `_MemoryStream` with.
+    """
+    try:
+        return asyncio.get_event_loop()
+    except RuntimeError:
+        # No loop is set on this thread. That is always the case in a sync
+        # context on Python 3.14+, where `get_event_loop()` no longer creates
+        # one implicitly, and on any version in a non-main thread.
+        #
+        # This loop is ours: unlike one returned by `get_event_loop()`, no
+        # caller owns it, so close it once `owner` dies. Otherwise the loop's
+        # `__del__` emits `ResourceWarning: unclosed event loop` at whatever
+        # nondeterministic point GC happens to run.
+        loop = asyncio.new_event_loop()
+        weakref.finalize(owner, _close_loop_quietly, loop)
+        return loop
+
+
+class _SendStream(Generic[T]):
+    def __init__(
+        self, reader_loop: AbstractEventLoop, queue: Queue[Any], done: object
+    ) -> None:
+        """Create a writer for the queue and done object.
+
+        Args:
+            reader_loop: The event loop to use for the writer.
+
+                This loop will be used to schedule the writes to the queue.
+            queue: The queue to write to.
+
+                This is an asyncio queue.
+            done: Special sentinel object to indicate that the writer is done.
+        """
+        self._reader_loop = reader_loop
+        self._queue = queue
+        self._done = done
+
+    async def send(self, item: T) -> None:
+        """Schedule the item to be written to the queue using the original loop.
+
+        This is a coroutine that can be awaited.
+
+        Args:
+            item: The item to write to the queue.
+        """
+        return self.send_nowait(item)
+
+    def send_nowait(self, item: T) -> None:
+        """Schedule the item to be written to the queue using the original loop.
+
+        This is a non-blocking call.
+
+        Args:
+            item: The item to write to the queue.
+
+        Raises:
+            RuntimeError: If the event loop is already closed when trying to write to
+                the queue.
+        """
+        try:
+            self._reader_loop.call_soon_threadsafe(self._queue.put_nowait, item)
+        except RuntimeError:
+            if not self._reader_loop.is_closed():
+                raise  # Raise the exception if the loop is not closed
+
+    async def aclose(self) -> None:
+        """Async schedule the done object write the queue using the original loop."""
+        return self.close()
+
+    def close(self) -> None:
+        """Schedule the done object write the queue using the original loop.
+
+        This is a non-blocking call.
+
+        Raises:
+            RuntimeError: If the event loop is already closed when trying to write to
+                the queue.
+        """
+        try:
+            self._reader_loop.call_soon_threadsafe(self._queue.put_nowait, self._done)
+        except RuntimeError:
+            if not self._reader_loop.is_closed():
+                raise  # Raise the exception if the loop is not closed
+
+
+class _ReceiveStream(Generic[T]):
+    def __init__(self, queue: Queue[Any], done: object) -> None:
+        """Create a reader for the queue and done object.
+
+        This reader should be used in the same loop as the loop that was passed to the
+        channel.
+        """
+        self._queue = queue
+        self._done = done
+        self._is_closed = False
+
+    async def __aiter__(self) -> AsyncIterator[T]:
+        while True:
+            item = await self._queue.get()
+            if item is self._done:
+                self._is_closed = True
+                break
+            yield item
+
+
+class _MemoryStream(Generic[T]):
+    """Stream data from a writer to a reader even if they are in different threads.
+
+    Uses asyncio queues to communicate between two co-routines. This implementation
+    should work even if the writer and reader co-routines belong to two different event
+    loops (e.g. one running from an event loop in the main thread and the other running
+    in an event loop in a background thread).
+
+    This implementation is meant to be used with a single writer and a single reader.
+
+    This is an internal implementation to Coco. Do not use it directly.
+    """
+
+    def __init__(self, loop: AbstractEventLoop) -> None:
+        """Create a channel for the given loop.
+
+        Args:
+            loop: The event loop to use for the channel.
+
+                The reader is assumed to be running in the same loop as the one passed
+                to this constructor. This will NOT be validated at run time.
+        """
+        self._loop = loop
+        self._queue = asyncio.Queue[Any](maxsize=0)
+        self._done = object()
+
+    def get_send_stream(self) -> _SendStream[T]:
+        """Get a writer for the channel.
+
+        Returns:
+            The writer for the channel.
+        """
+        return _SendStream[T](
+            reader_loop=self._loop, queue=self._queue, done=self._done
+        )
+
+    def get_receive_stream(self) -> _ReceiveStream[T]:
+        """Get a reader for the channel.
+
+        Returns:
+            The reader for the channel.
+        """
+        return _ReceiveStream[T](queue=self._queue, done=self._done)

coco_core/tracers/root_listeners.py

@@ -0,0 +1,130 @@
+"""Tracers that call listeners."""
+
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING
+
+from coco_core.runnables.config import (
+    RunnableConfig,
+    acall_func_with_variable_args,
+    call_func_with_variable_args,
+)
+from coco_core.tracers.base import AsyncBaseTracer, BaseTracer
+from coco_core.tracers.schemas import Run
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+Listener = Callable[[Run], None] | Callable[[Run, RunnableConfig], None]
+AsyncListener = (
+    Callable[[Run], Awaitable[None]] | Callable[[Run, RunnableConfig], Awaitable[None]]
+)
+
+
+class RootListenersTracer(BaseTracer):
+    """Tracer that calls listeners on run start, end, and error."""
+
+    log_missing_parent = False
+    """Whether to log a warning if the parent is missing."""
+
+    def __init__(
+        self,
+        *,
+        config: RunnableConfig,
+        on_start: Listener | None,
+        on_end: Listener | None,
+        on_error: Listener | None,
+    ) -> None:
+        """Initialize the tracer.
+
+        Args:
+            config: The runnable config.
+            on_start: The listener to call on run start.
+            on_end: The listener to call on run end.
+            on_error: The listener to call on run error
+        """
+        super().__init__(_schema_format="original+chat")
+
+        self.config = config
+        self._arg_on_start = on_start
+        self._arg_on_end = on_end
+        self._arg_on_error = on_error
+        self.root_id: UUID | None = None
+
+    def _persist_run(self, run: Run) -> None:
+        # This is a legacy method only called once for an entire run tree
+        # therefore not useful here
+        pass
+
+    def _on_run_create(self, run: Run) -> None:
+        if self.root_id is not None:
+            return
+
+        self.root_id = run.id
+
+        if self._arg_on_start is not None:
+            call_func_with_variable_args(self._arg_on_start, run, self.config)
+
+    def _on_run_update(self, run: Run) -> None:
+        if run.id != self.root_id:
+            return
+
+        if run.error is None:
+            if self._arg_on_end is not None:
+                call_func_with_variable_args(self._arg_on_end, run, self.config)
+        elif self._arg_on_error is not None:
+            call_func_with_variable_args(self._arg_on_error, run, self.config)
+
+
+class AsyncRootListenersTracer(AsyncBaseTracer):
+    """Async tracer that calls listeners on run start, end, and error."""
+
+    log_missing_parent = False
+    """Whether to log a warning if the parent is missing."""
+
+    def __init__(
+        self,
+        *,
+        config: RunnableConfig,
+        on_start: AsyncListener | None,
+        on_end: AsyncListener | None,
+        on_error: AsyncListener | None,
+    ) -> None:
+        """Initialize the tracer.
+
+        Args:
+            config: The runnable config.
+            on_start: The listener to call on run start.
+            on_end: The listener to call on run end.
+            on_error: The listener to call on run error
+        """
+        super().__init__(_schema_format="original+chat")
+
+        self.config = config
+        self._arg_on_start = on_start
+        self._arg_on_end = on_end
+        self._arg_on_error = on_error
+        self.root_id: UUID | None = None
+
+    async def _persist_run(self, run: Run) -> None:
+        # This is a legacy method only called once for an entire run tree
+        # therefore not useful here
+        pass
+
+    async def _on_run_create(self, run: Run) -> None:
+        if self.root_id is not None:
+            return
+
+        self.root_id = run.id
+
+        if self._arg_on_start is not None:
+            await acall_func_with_variable_args(self._arg_on_start, run, self.config)
+
+    async def _on_run_update(self, run: Run) -> None:
+        if run.id != self.root_id:
+            return
+
+        if run.error is None:
+            if self._arg_on_end is not None:
+                await acall_func_with_variable_args(self._arg_on_end, run, self.config)
+        elif self._arg_on_error is not None:
+            await acall_func_with_variable_args(self._arg_on_error, run, self.config)

coco_core/tracers/run_collector.py

@@ -0,0 +1,40 @@
+"""A tracer that collects all nested runs in a list."""
+
+from typing import Any
+from uuid import UUID
+
+from coco_core.tracers._compat import run_copy
+from coco_core.tracers.base import BaseTracer
+from coco_core.tracers.schemas import Run
+
+
+class RunCollectorCallbackHandler(BaseTracer):
+    """Tracer that collects all nested runs in a list.
+
+    This tracer is useful for inspection and evaluation purposes.
+    """
+
+    name: str = "run-collector_callback_handler"
+
+    def __init__(self, example_id: UUID | str | None = None, **kwargs: Any) -> None:
+        """Initialize the `RunCollectorCallbackHandler`.
+
+        Args:
+            example_id: The ID of the example being traced.
+            **kwargs: Additional keyword arguments.
+        """
+        super().__init__(**kwargs)
+        self.example_id = (
+            UUID(example_id) if isinstance(example_id, str) else example_id
+        )
+        self.traced_runs: list[Run] = []
+
+    def _persist_run(self, run: Run) -> None:
+        """Persist a run by adding it to the `traced_runs` list.
+
+        Args:
+            run: The run to be persisted.
+        """
+        run_ = run_copy(run)
+        run_.reference_example_id = self.example_id
+        self.traced_runs.append(run_)

coco_core/tracers/schemas.py

@@ -0,0 +1,14 @@
+"""Schemas for tracers."""
+
+from __future__ import annotations
+
+from coco_core.tracers._context import RunTree
+
+# Begin V2 API Schemas
+
+
+Run = RunTree  # For backwards compatibility
+
+__all__ = [
+    "Run",
+]

coco_core/tracers/stdout.py

@@ -0,0 +1,205 @@
+"""Tracers that print to the console."""
+
+import json
+from collections.abc import Callable
+from typing import Any
+
+from coco_core.tracers.base import BaseTracer
+from coco_core.tracers.schemas import Run
+from coco_core.utils.input import get_bolded_text, get_colored_text
+
+MILLISECONDS_IN_SECOND = 1000
+
+
+def try_json_stringify(obj: Any, fallback: str) -> str:
+    """Try to stringify an object to JSON.
+
+    Args:
+        obj: Object to stringify.
+        fallback: Fallback string to return if the object cannot be stringified.
+
+    Returns:
+        A JSON string if the object can be stringified, otherwise the fallback string.
+    """
+    try:
+        return json.dumps(obj, indent=2, ensure_ascii=False)
+    except Exception:
+        return fallback
+
+
+def elapsed(run: Any) -> str:
+    """Get the elapsed time of a run.
+
+    Args:
+        run: any object with a `start_time` and `end_time` attribute.
+
+    Returns:
+        A string with the elapsed time in seconds or milliseconds if time is less than a
+            second.
+
+    """
+    elapsed_time = run.end_time - run.start_time
+    seconds = elapsed_time.total_seconds()
+    if seconds < 1:
+        return f"{seconds * MILLISECONDS_IN_SECOND:.0f}ms"
+    return f"{seconds:.2f}s"
+
+
+class FunctionCallbackHandler(BaseTracer):
+    """Tracer that calls a function with a single str parameter."""
+
+    name: str = "function_callback_handler"
+    """The name of the tracer.
+
+    This is used to identify the tracer in the logs.
+    """
+
+    def __init__(self, function: Callable[[str], None], **kwargs: Any) -> None:
+        """Create a `FunctionCallbackHandler`.
+
+        Args:
+            function: The callback function to call.
+        """
+        super().__init__(**kwargs)
+        self.function_callback = function
+
+    def _persist_run(self, run: Run) -> None:
+        pass
+
+    def get_parents(self, run: Run) -> list[Run]:
+        """Get the parents of a run.
+
+        Args:
+            run: The run to get the parents of.
+
+        Returns:
+            A list of parent runs.
+        """
+        parents = []
+        current_run = run
+        while current_run.parent_run_id:
+            parent = self.run_map.get(str(current_run.parent_run_id))
+            if parent:
+                parents.append(parent)
+                current_run = parent
+            else:
+                break
+        return parents
+
+    def get_breadcrumbs(self, run: Run) -> str:
+        """Get the breadcrumbs of a run.
+
+        Args:
+            run: The run to get the breadcrumbs of.
+
+        Returns:
+            A string with the breadcrumbs of the run.
+        """
+        parents = self.get_parents(run)[::-1]
+        return " > ".join(
+            f"{parent.run_type}:{parent.name}"
+            for i, parent in enumerate([*parents, run])
+        )
+
+    # logging methods
+    def _on_chain_start(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        run_type = run.run_type.capitalize()
+        self.function_callback(
+            f"{get_colored_text('[chain/start]', color='green')} "
+            + get_bolded_text(f"[{crumbs}] Entering {run_type} run with input:\n")
+            + f"{try_json_stringify(run.inputs, '[inputs]')}"
+        )
+
+    def _on_chain_end(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        run_type = run.run_type.capitalize()
+        self.function_callback(
+            f"{get_colored_text('[chain/end]', color='blue')} "
+            + get_bolded_text(
+                f"[{crumbs}] [{elapsed(run)}] Exiting {run_type} run with output:\n"
+            )
+            + f"{try_json_stringify(run.outputs, '[outputs]')}"
+        )
+
+    def _on_chain_error(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        run_type = run.run_type.capitalize()
+        self.function_callback(
+            f"{get_colored_text('[chain/error]', color='red')} "
+            + get_bolded_text(
+                f"[{crumbs}] [{elapsed(run)}] {run_type} run errored with error:\n"
+            )
+            + f"{try_json_stringify(run.error, '[error]')}"
+        )
+
+    def _on_llm_start(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        inputs = (
+            {"prompts": [p.strip() for p in run.inputs["prompts"]]}
+            if "prompts" in run.inputs
+            else run.inputs
+        )
+        self.function_callback(
+            f"{get_colored_text('[llm/start]', color='green')} "
+            + get_bolded_text(f"[{crumbs}] Entering LLM run with input:\n")
+            + f"{try_json_stringify(inputs, '[inputs]')}"
+        )
+
+    def _on_llm_end(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        self.function_callback(
+            f"{get_colored_text('[llm/end]', color='blue')} "
+            + get_bolded_text(
+                f"[{crumbs}] [{elapsed(run)}] Exiting LLM run with output:\n"
+            )
+            + f"{try_json_stringify(run.outputs, '[response]')}"
+        )
+
+    def _on_llm_error(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        self.function_callback(
+            f"{get_colored_text('[llm/error]', color='red')} "
+            + get_bolded_text(
+                f"[{crumbs}] [{elapsed(run)}] LLM run errored with error:\n"
+            )
+            + f"{try_json_stringify(run.error, '[error]')}"
+        )
+
+    def _on_tool_start(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        self.function_callback(
+            f"{get_colored_text('[tool/start]', color='green')} "
+            + get_bolded_text(f"[{crumbs}] Entering Tool run with input:\n")
+            + f'"{run.inputs["input"].strip()}"'
+        )
+
+    def _on_tool_end(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        if run.outputs:
+            self.function_callback(
+                f"{get_colored_text('[tool/end]', color='blue')} "
+                + get_bolded_text(
+                    f"[{crumbs}] [{elapsed(run)}] Exiting Tool run with output:\n"
+                )
+                + f'"{str(run.outputs["output"]).strip()}"'
+            )
+
+    def _on_tool_error(self, run: Run) -> None:
+        crumbs = self.get_breadcrumbs(run)
+        self.function_callback(
+            f"{get_colored_text('[tool/error]', color='red')} "
+            + get_bolded_text(f"[{crumbs}] [{elapsed(run)}] ")
+            + f"Tool run errored with error:\n"
+            f"{run.error}"
+        )
+
+
+class ConsoleCallbackHandler(FunctionCallbackHandler):
+    """Tracer that prints to the console."""
+
+    name: str = "console_callback_handler"
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Create a ConsoleCallbackHandler."""
+        super().__init__(function=print, **kwargs)

coco_core/utils/__init__.py

@@ -0,0 +1,111 @@
+"""Utility functions for Coco.
+
+These functions do not depend on any other Coco module.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    # for type checking and IDE support, we include the imports here
+    # but we don't want to eagerly import them at runtime
+    from coco_core.utils import image
+    from coco_core.utils.aiter import abatch_iterate
+    from coco_core.utils.env import get_from_dict_or_env, get_from_env
+    from coco_core.utils.formatting import StrictFormatter, formatter
+    from coco_core.utils.input import (
+        get_bolded_text,
+        get_color_mapping,
+        get_colored_text,
+        print_text,
+    )
+    from coco_core.utils.iter import batch_iterate
+    from coco_core.utils.pydantic import pre_init
+    from coco_core.utils.strings import (
+        comma_list,
+        sanitize_for_postgres,
+        stringify_dict,
+        stringify_value,
+    )
+    from coco_core.utils.utils import (
+        build_extra_kwargs,
+        check_package_version,
+        convert_to_secret_str,
+        from_env,
+        get_pydantic_field_names,
+        guard_import,
+        mock_now,
+        raise_for_status_with_text,
+        secret_from_env,
+        xor_args,
+    )
+
+__all__ = (
+    "StrictFormatter",
+    "abatch_iterate",
+    "batch_iterate",
+    "build_extra_kwargs",
+    "check_package_version",
+    "comma_list",
+    "convert_to_secret_str",
+    "formatter",
+    "from_env",
+    "get_bolded_text",
+    "get_color_mapping",
+    "get_colored_text",
+    "get_from_dict_or_env",
+    "get_from_env",
+    "get_pydantic_field_names",
+    "guard_import",
+    "image",
+    "mock_now",
+    "pre_init",
+    "print_text",
+    "raise_for_status_with_text",
+    "sanitize_for_postgres",
+    "secret_from_env",
+    "stringify_dict",
+    "stringify_value",
+    "xor_args",
+)
+
+_dynamic_imports = {
+    "image": "__module__",
+    "abatch_iterate": "aiter",
+    "get_from_dict_or_env": "env",
+    "get_from_env": "env",
+    "StrictFormatter": "formatting",
+    "formatter": "formatting",
+    "get_bolded_text": "input",
+    "get_color_mapping": "input",
+    "get_colored_text": "input",
+    "print_text": "input",
+    "batch_iterate": "iter",
+    "pre_init": "pydantic",
+    "comma_list": "strings",
+    "sanitize_for_postgres": "strings",
+    "stringify_dict": "strings",
+    "stringify_value": "strings",
+    "build_extra_kwargs": "utils",
+    "check_package_version": "utils",
+    "convert_to_secret_str": "utils",
+    "from_env": "utils",
+    "get_pydantic_field_names": "utils",
+    "guard_import": "utils",
+    "mock_now": "utils",
+    "secret_from_env": "utils",
+    "xor_args": "utils",
+    "raise_for_status_with_text": "utils",
+}
+
+
+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__)