nat.profiler.decorators.framework_wrapper

Attributes

logger
_library_instrumented
callback_handler_var

Functions

set_framework_profiler_handler(...)	Decorator that wraps an async context manager function to set up framework-specific profiling.
chain_wrapped_build_fn(...)	Convert an original build function into an async context manager that

Module Contents

logger

_library_instrumented

callback_handler_var: contextvars.ContextVar[Any | None]

set_framework_profiler_handler(

workflow_llms: dict | None = None,

frameworks: list[nat.builder.framework_enum.LLMFrameworkEnum] | None = None,

) → collections.abc.Callable[[collections.abc.Callable[Ellipsis, contextlib.AbstractAsyncContextManager[Any]]], collections.abc.Callable[Ellipsis, contextlib.AbstractAsyncContextManager[Any]]]

Decorator that wraps an async context manager function to set up framework-specific profiling.

Args:

workflow_llms (dict | None): A dictionary of workflow LLM configurations. frameworks (list[LLMFrameworkEnum] | None): A list of LLM frameworks used in the workflow functions.

Returns:

Callable[[Callable[…, AsyncContextManager[Any]]], Callable[…, AsyncContextManager[Any]]]: A decorator that wraps the original function with profiling setup.

chain_wrapped_build_fn(

original_build_fn: collections.abc.Callable[Ellipsis, contextlib.AbstractAsyncContextManager],

workflow_llms: dict,

function_frameworks: list[nat.builder.framework_enum.LLMFrameworkEnum],

) → collections.abc.Callable[Ellipsis, contextlib.AbstractAsyncContextManager]

Convert an original build function into an async context manager that wraps it with a single call to set_framework_profiler_handler, passing all frameworks at once.

Args:

original_build_fn (Callable[…, AsyncContextManager]): The original build function to wrap. workflow_llms (dict): A dictionary of workflow LLM configurations. function_frameworks (list[LLMFrameworkEnum]): A list of LLM frameworks used in the workflow functions.

Returns:

Callable[…, AsyncContextManager]: The wrapped build function.