microsoft--agent-lightning/agentlightning/algorithm/decorator.py

# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import functools
import inspect
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    Dict,
    Generic,
    Literal,
    Optional,
    Protocol,
    TypeVar,
    Union,
    cast,
    overload,
)

from agentlightning.adapter import TraceAdapter
from agentlightning.store.base import LightningStore
from agentlightning.types import Dataset, NamedResources

if TYPE_CHECKING:
    from agentlightning.llm_proxy import LLMProxy

from .base import Algorithm

# Algorithm function signature types
# We've missed a lot of combinations here.
# Let's add them in future.


class AlgorithmFuncSyncFull(Protocol):
    def __call__(
        self,
        *,
        store: LightningStore,
        train_dataset: Optional[Dataset[Any]],
        val_dataset: Optional[Dataset[Any]],
        llm_proxy: Optional[LLMProxy],
        adapter: Optional[TraceAdapter[Any]],
        initial_resources: Optional[NamedResources],
    ) -> None: ...


class AlgorithmFuncSyncOnlyStore(Protocol):
    def __call__(self, *, store: LightningStore) -> None: ...


class AlgorithmFuncSyncOnlyDataset(Protocol):
    def __call__(self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None: ...


class AlgorithmFuncAsyncFull(Protocol):
    def __call__(
        self,
        *,
        store: LightningStore,
        train_dataset: Optional[Dataset[Any]],
        val_dataset: Optional[Dataset[Any]],
        llm_proxy: Optional[LLMProxy],
        adapter: Optional[TraceAdapter[Any]],
        initial_resources: Optional[NamedResources],
    ) -> Awaitable[None]: ...


class AlgorithmFuncAsyncOnlyStore(Protocol):
    def __call__(self, *, store: LightningStore) -> Awaitable[None]: ...


class AlgorithmFuncAsyncOnlyDataset(Protocol):
    def __call__(
        self, *, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]
    ) -> Awaitable[None]: ...


AlgorithmFuncAsync = Union[AlgorithmFuncAsyncOnlyStore, AlgorithmFuncAsyncOnlyDataset, AlgorithmFuncAsyncFull]

AlgorithmFuncSync = Union[AlgorithmFuncSyncOnlyStore, AlgorithmFuncSyncOnlyDataset, AlgorithmFuncSyncFull]


class AlgorithmFuncSyncFallback(Protocol):
    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...


class AlgorithmFuncAsyncFallback(Protocol):
    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]: ...


AlgorithmFuncSyncLike = Union[AlgorithmFuncSync, AlgorithmFuncSyncFallback]
AlgorithmFuncAsyncLike = Union[AlgorithmFuncAsync, AlgorithmFuncAsyncFallback]

AlgorithmFunc = Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]


AsyncFlag = Literal[True, False]
AF = TypeVar("AF", bound=AsyncFlag)


class FunctionalAlgorithm(Algorithm, Generic[AF]):
    """An algorithm wrapper built from a callable implementation.

    Functional algorithms let you provide an ordinary function instead of
    subclassing [`Algorithm`][agentlightning.Algorithm]. The wrapper inspects
    the callable signature to supply optional dependencies
    such as the store, adapter, and LLM proxy.
    """

    @overload
    def __init__(self: "FunctionalAlgorithm[Literal[False]]", algorithm_func: AlgorithmFuncSyncLike) -> None: ...

    @overload
    def __init__(self: "FunctionalAlgorithm[Literal[True]]", algorithm_func: AlgorithmFuncAsyncLike) -> None: ...

    def __init__(self, algorithm_func: Union[AlgorithmFuncSyncLike, AlgorithmFuncAsyncLike]) -> None:
        """Wrap a function that implements algorithm behaviour.

        Args:
            algorithm_func: Sync or async callable implementing the algorithm
                contract. Arguments are detected automatically based on the
                function signature.
        """
        super().__init__()
        self._algorithm_func = algorithm_func
        self._sig = inspect.signature(algorithm_func)
        self._is_async = inspect.iscoroutinefunction(algorithm_func)

        # Copy function metadata to preserve type hints and other attributes
        functools.update_wrapper(self, algorithm_func)  # type: ignore

    def is_async(self) -> bool:
        return self._is_async

    @overload
    def run(
        self: "FunctionalAlgorithm[Literal[False]]",
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> None: ...

    @overload
    def run(
        self: "FunctionalAlgorithm[Literal[True]]",
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> Awaitable[None]: ...

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        return self._algorithm_func(*args, **kwargs)  # type: ignore

    def run(
        self,
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> Union[None, Awaitable[None]]:
        """Execute the wrapped function with injected dependencies.

        Args:
            train_dataset: Optional training dataset passed through when the
                callable declares a `train_dataset` parameter.
            val_dataset: Optional validation dataset passed through when the
                callable declares a `val_dataset` parameter.

        Returns:
            None for sync callables or an awaitable when the callable is async.

        Raises:
            TypeError: If a dataset is provided but the function signature does
                not accept the corresponding argument.
        """
        kwargs: Dict[str, Any] = {}
        if "store" in self._sig.parameters:
            kwargs["store"] = self.get_store()
        if "adapter" in self._sig.parameters:
            kwargs["adapter"] = self.get_adapter()
        if "llm_proxy" in self._sig.parameters:
            kwargs["llm_proxy"] = self.get_llm_proxy()
        if "initial_resources" in self._sig.parameters:
            kwargs["initial_resources"] = self.get_initial_resources()
        if "train_dataset" in self._sig.parameters:
            kwargs["train_dataset"] = train_dataset
        elif train_dataset is not None:
            raise TypeError(
                f"train_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
            )
        if "val_dataset" in self._sig.parameters:
            kwargs["val_dataset"] = val_dataset
        elif val_dataset is not None:
            raise TypeError(
                f"val_dataset is provided but not supported by the algorithm function: {self._algorithm_func}"
            )
        # both sync and async functions can be called with the same signature
        result = self._algorithm_func(**kwargs)  # type: ignore[misc]
        if self._is_async:
            return cast(Awaitable[None], result)
        return None


@overload
def algo(func: AlgorithmFuncAsync) -> FunctionalAlgorithm[Literal[True]]: ...


@overload
def algo(func: AlgorithmFuncAsyncFallback) -> FunctionalAlgorithm[Any]: ...


@overload
def algo(func: AlgorithmFuncSync) -> FunctionalAlgorithm[Literal[False]]: ...


@overload
def algo(func: AlgorithmFuncSyncFallback) -> FunctionalAlgorithm[Any]: ...


def algo(
    func: Union[
        AlgorithmFuncSync,
        AlgorithmFuncAsync,
        AlgorithmFuncSyncFallback,
        AlgorithmFuncAsyncFallback,
    ],
) -> Union[FunctionalAlgorithm[Literal[False]], FunctionalAlgorithm[Literal[True]]]:
    """Convert a callable into a [`FunctionalAlgorithm`][agentlightning.algorithm.decorator.FunctionalAlgorithm].

    The decorator inspects the callable signature to decide which dependencies
    to inject at runtime, enabling concise algorithm definitions that still
    leverage the full training runtime.

    Args:
        func: Function implementing the algorithm logic. May be synchronous or
            asynchronous. The function can expect all of, or a subset of the following parameters:

            - `store`: [`LightningStore`][agentlightning.store.base.LightningStore],
            - `train_dataset`: [`Dataset`][agentlightning.Dataset],
            - `val_dataset`: [`Dataset`][agentlightning.Dataset],
            - `llm_proxy`: [`LLMProxy`][agentlightning.LLMProxy],
            - `adapter`: [`TraceAdapter`][agentlightning.TraceAdapter],
            - `initial_resources`: [`NamedResources`][agentlightning.NamedResources],

            If the function does not expect a parameter, the wrapper will not inject it into the call.
            Using `*args` and `**kwargs` will not work and no parameters will be injected.

    Returns:
        FunctionalAlgorithm that proxies the callable while exposing the
        `Algorithm` interface.

    Examples:
        ```python
        from agentlightning.algorithm.decorator import algo

        @algo
        def batching_algorithm(*, store, train_dataset, val_dataset):
            for sample in train_dataset:
                store.enqueue_rollout(input=sample, mode="train")

        @algo
        async def async_algorithm(*, store, train_dataset=None, val_dataset=None):
            await store.enqueue_rollout(input={"prompt": "hello"}, mode="train")
        ```
    """
    return FunctionalAlgorithm(func)