decorators

Class info

🛈 DocStrings

Decorators for agent injection and execution.

with_nodes

with_nodes(
    pool: AgentPool,
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]

with_nodes(
    pool: AgentPool, func: Callable[P, Awaitable[T]]
) -> Callable[P, Awaitable[T]]

with_nodes(
    pool: AgentPool, func: Callable[P, Awaitable[T]] | None = None
) -> (
    Callable[P, Awaitable[T]]
    | Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]
)

Inject nodes into function parameters.

Source code in src/llmling_agent/running/decorators.py

def with_nodes[T, **P](
    pool: AgentPool,
    func: Callable[P, Awaitable[T]] | None = None,
) -> (
    Callable[P, Awaitable[T]]
    | Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]
):
    """Inject nodes into function parameters."""

    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
        @wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            # Convert args to kwargs for injection check
            sig = inspect.signature(func)
            bound_args = sig.bind_partial(*args)
            all_kwargs = {**bound_args.arguments, **kwargs}

            # Get needed nodes
            nodes = inject_nodes(func, pool, all_kwargs)

            # Create kwargs with nodes first, then other args
            final_kwargs = {**nodes, **kwargs}

            # Convert back to args/kwargs using signature
            bound = sig.bind(**final_kwargs)
            bound.apply_defaults()

            # Call with proper args/kwargs
            return await func(*bound.args, **bound.kwargs)

        return wrapper

    return decorator(func) if func else decorator