Source code for langchain_core.runnables.schema

"""Module contains typedefs that are used with Runnables."""

from __future__ import annotations

from typing import Any, Dict, List, Literal, Sequence, Union

from typing_extensions import NotRequired, TypedDict


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.
    """
    output: Any
    """The output of the Runnable that generated the event.
    
    Outputs will only be available at the *END* of the Runnable.
    
    For most Runnables, this field can be inferred from the `chunk` field,
    though there might be some exceptions for special cased Runnables (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.
    """


class BaseStreamEvent(TypedDict):
    """Streaming event.

    Schema of a streaming event which is produced from the astream_events method.

    Example:

        .. code-block:: python

            from langchain_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 Runnables 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 Runnables.
    
    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 LangChain 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.

    .. versionadded:: 0.2.15
    """

    # 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 = Union[StandardStreamEvent, CustomStreamEvent]