More decorator support for LCEL chains

[benobytes]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

It would be nice if I could use more decorators in my code as (personal opinion) I believe it becomes more readable.

Motivation

I really like the concept of using a decorator for chains instead of constantly defining RunnableLambda everywhere, but it feels underdeveloped.

Proposal (If applicable)

Here is my POC

"""Useful decorators for configuring runnables."""

from langchain_core.callbacks import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)
from langchain_core.runnables import (
    Runnable,
    RunnableConfig,
    RunnableLambda,
)
from langchain_core.tracers.schemas import Run
from typing_extensions import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Iterator,
    Optional,
    Union,
)

Input = Any
Output = Any
Listener = Union[Callable[[Run], None], Callable[[Run, RunnableConfig], None]]

# I included this first part because I was experiencing IDE type issues with the old @chain decorator when I had a RunnableConfig in my function signature. 
def chain(
    func: Union[
        Union[
            Callable[[Input], Output],
            Callable[[Input], Iterator[Output]],
            Callable[[Input, RunnableConfig], Output],
            Callable[[Input, CallbackManagerForChainRun], Output],
            Callable[[Input, CallbackManagerForChainRun, RunnableConfig], Output],
        ],
        Union[
            Callable[[Input], Awaitable[Output]],
            Callable[[Input], AsyncIterator[Output]],
            Callable[[Input, RunnableConfig], Awaitable[Output]],
            Callable[[Input, AsyncCallbackManagerForChainRun], Awaitable[Output]],
            Callable[
                [Input, AsyncCallbackManagerForChainRun, RunnableConfig],
                Awaitable[Output],
            ],
        ],
    ],
) -> Runnable[Input, Output]:
    """Decorate a function to make it a Runnable.
    Sets the name of the runnable to the name of the function.
    Any runnables called by the function will be traced as dependencies.

    Args:
        func: A callable.

    Returns:
        A Runnable.
    """
    return RunnableLambda(func)


def with_config(
    config: Optional[RunnableConfig] = None,
    **kwargs: Any,
) -> Callable[..., Runnable[Input, Output]]:
    """Bind a Runnable to a RunnableConfig.

    Args:
        runnable: A Runnable.

    Returns:
        A RunnableBinding.
    """

    def _with_config(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_config(config=config, **kwargs)

    return _with_config


def with_listeners(
    *,
    on_start: Optional[Listener] = None,
    on_end: Optional[Listener] = None,
    on_error: Optional[Listener] = None,
) -> Callable[..., Runnable[Input, Output]]:
    """Bind lifecycle listeners to a Runnable, returning a new Runnable.

    on_start: Called before the runnable starts running, with the Run object.
    on_end: Called after the runnable finishes running, with the Run object.
    on_error: Called if the runnable throws an error, with the Run object.

    The Run object contains information about the run, including its id,
    type, input, output, error, start_time, end_time, and any tags or metadata
    added to the run.
    """

    def _with_listeners(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_listeners(
            on_start=on_start, on_end=on_end, on_error=on_error
        )

    return _with_listeners


def with_types(
    input_type: Optional[type] = None,
    output_type: Optional[type] = None,
) -> Callable[..., Runnable[Input, Output]]:
    """Bind input and output types to a Runnable, returning a new Runnable."""

    def _with_types(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_types(input_type=input_type, output_type=output_type)

    return _with_types


def with_retry(
    retry_if_exception_type: tuple[type[BaseException], ...] = (Exception,),
    wait_exponential_jitter: bool = True,
    stop_after_attempt: int = 3,
) -> Callable[..., Runnable[Input, Output]]:
    """Create a new Runnable that retries the original runnable on exceptions."""

    def _with_retry(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_retry(
            retry_if_exception_type=retry_if_exception_type,
            wait_exponential_jitter=wait_exponential_jitter,
            stop_after_attempt=stop_after_attempt,
        )

    return _with_retry


def with_fallbacks(
    fallbacks: list[Runnable[Input, Output]],
    *,
    exceptions_to_handle: tuple[type[BaseException], ...] = (Exception,),
    exception_key: Optional[str] = None,
) -> Callable[..., Runnable[Input, Output]]:
    """Add fallbacks to a runnable, returning a new Runnable."""

    def _with_fallbacks(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_fallbacks(
            fallbacks=fallbacks,
            exceptions_to_handle=exceptions_to_handle,
            exception_key=exception_key,
        )

    return _with_fallbacks


def on_chain_start(listener: Listener) -> Callable[..., Runnable[Input, Output]]:
    """Bind a listener to the start of a Runnable, returning a new Runnable."""

    def _on_start(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_listeners(on_start=listener)

    return _on_start


def on_chain_end(listener: Listener) -> Callable[..., Runnable[Input, Output]]:
    """Bind a listener to the end of a Runnable, returning a new Runnable."""

    def _on_end(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_listeners(on_end=listener)

    return _on_end


def on_chain_error(listener: Listener) -> Callable[..., Runnable[Input, Output]]:
    """Bind a listener to the error of a Runnable, returning a new Runnable."""

    def _on_error(runnable: Runnable[Input, Output]) -> Runnable[Input, Output]:
        return runnable.with_listeners(on_error=listener)

    return _on_error