coco-framework 0.1.0__py3-none-any.whl

coco_core/runnables/schema.py

@@ -0,0 +1,188 @@
+"""Module contains typedefs that are used with `Runnable` objects."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from typing_extensions import NotRequired, TypedDict
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+class EventData(TypedDict, total=False):
+    """Data associated with a streaming event."""
+
+    input: Any
+    """The input passed to the `Runnable` that generated the event.
+
+    Inputs will sometimes be available at the *START* of the `Runnable`, and
+    sometimes at the *END* of the `Runnable`.
+
+    If a `Runnable` is able to stream its inputs, then its input by definition
+    won't be known until the *END* of the `Runnable` when it has finished streaming
+    its inputs.
+    """
+    error: NotRequired[BaseException]
+    """The error that occurred during the execution of the `Runnable`.
+
+    This field is only available if the `Runnable` raised an exception.
+
+    !!! version-added "Added in `trustcoco-core` 1.0.0"
+    """
+    output: Any
+    """The output of the `Runnable` that generated the event.
+
+    Outputs will only be available at the *END* of the `Runnable`.
+
+    For most `Runnable` objects, this field can be inferred from the `chunk` field,
+    though there might be some exceptions for special a cased `Runnable` (e.g., like
+    chat models), which may return more information.
+    """
+    chunk: Any
+    """A streaming chunk from the output that generated the event.
+
+    chunks support addition in general, and adding them up should result
+    in the output of the `Runnable` that generated the event.
+    """
+    tool_call_id: NotRequired[str | None]
+    """The tool call ID associated with the tool execution.
+
+    This field is available for the `on_tool_error` event and can be used to
+    link errors to specific tool calls in stateless agent implementations.
+    """
+
+
+class BaseStreamEvent(TypedDict):
+    """Streaming event.
+
+    Schema of a streaming event which is produced from the `astream_events` method.
+
+    Example:
+        ```python
+        from coco_core.runnables import RunnableLambda
+
+
+        async def reverse(s: str) -> str:
+            return s[::-1]
+
+
+        chain = RunnableLambda(func=reverse)
+
+        events = [event async for event in chain.astream_events("hello")]
+
+        # Will produce the following events
+        # (where some fields have been omitted for brevity):
+        [
+            {
+                "data": {"input": "hello"},
+                "event": "on_chain_start",
+                "metadata": {},
+                "name": "reverse",
+                "tags": [],
+            },
+            {
+                "data": {"chunk": "olleh"},
+                "event": "on_chain_stream",
+                "metadata": {},
+                "name": "reverse",
+                "tags": [],
+            },
+            {
+                "data": {"output": "olleh"},
+                "event": "on_chain_end",
+                "metadata": {},
+                "name": "reverse",
+                "tags": [],
+            },
+        ]
+        ```
+    """
+
+    event: str
+    """Event names are of the format: `on_[runnable_type]_(start|stream|end)`.
+
+    Runnable types are one of:
+
+    - **llm** - used by non chat models
+    - **chat_model** - used by chat models
+    - **prompt** --  e.g., `ChatPromptTemplate`
+    - **tool** -- from tools defined via `@tool` decorator or inheriting
+        from `Tool`/`BaseTool`
+    - **chain** - most `Runnable` objects are of this type
+
+    Further, the events are categorized as one of:
+
+    - **start** - when the `Runnable` starts
+    - **stream** - when the `Runnable` is streaming
+    - **end* - when the `Runnable` ends
+
+    start, stream and end are associated with slightly different `data` payload.
+
+    Please see the documentation for `EventData` for more details.
+    """
+    run_id: str
+    """An randomly generated ID to keep track of the execution of the given `Runnable`.
+
+    Each child `Runnable` that gets invoked as part of the execution of a parent
+    `Runnable` is assigned its own unique ID.
+    """
+    tags: NotRequired[list[str]]
+    """Tags associated with the `Runnable` that generated this event.
+
+    Tags are always inherited from parent `Runnable` objects.
+
+    Tags can either be bound to a `Runnable` using `.with_config({"tags":  ["hello"]})`
+    or passed at run time using `.astream_events(..., {"tags": ["hello"]})`.
+    """
+    metadata: NotRequired[dict[str, Any]]
+    """Metadata associated with the `Runnable` that generated this event.
+
+    Metadata can either be bound to a `Runnable` using
+
+        `.with_config({"metadata": { "foo": "bar" }})`
+
+    or passed at run time using
+
+        `.astream_events(..., {"metadata": {"foo": "bar"}})`.
+    """
+
+    parent_ids: Sequence[str]
+    """A list of the parent IDs associated with this event.
+
+    Root Events will have an empty list.
+
+    For example, if a `Runnable` A calls `Runnable` B, then the event generated by
+    `Runnable` B will have `Runnable` A's ID in the `parent_ids` field.
+
+    The order of the parent IDs is from the root parent to the immediate parent.
+
+    Only supported as of v2 of the astream events API. v1 will return an empty list.
+    """
+
+
+class StandardStreamEvent(BaseStreamEvent):
+    """A standard stream event that follows Coco convention for event data."""
+
+    data: EventData
+    """Event data.
+
+    The contents of the event data depend on the event type.
+    """
+    name: str
+    """The name of the `Runnable` that generated the event."""
+
+
+class CustomStreamEvent(BaseStreamEvent):
+    """Custom stream event created by the user."""
+
+    # Overwrite the event field to be more specific.
+    event: Literal["on_custom_event"]  # type: ignore[misc]
+    """The event type."""
+    name: str
+    """User defined name for the event."""
+    data: Any
+    """The data associated with the event. Free form and can be anything."""
+
+
+StreamEvent = StandardStreamEvent | CustomStreamEvent