Skip to content

datarobot_genai.nat.tool

tool

nat_tool 

nat_tool(fn: Callable, name: str, description: str | None = None) -> Callable

Decorate a function as a NAT tool.

Source code in datarobot_genai/nat/tool.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def nat_tool(fn: Callable, name: str, description: str | None = None) -> Callable:
    """Decorate a function as a NAT tool."""

    class NatToolConfig(FunctionBaseConfig, name=name):  # type: ignore[call-arg]
        pass

    @register_function(
        config_type=NatToolConfig,
    )
    async def wrapper(config: NatToolConfig, builder: Builder) -> AsyncGenerator:
        # NAT expects a coroutine function, so we wrap sync functions in an async one.
        if not inspect.iscoroutinefunction(fn):
            fn_for_nat = _sync_to_async(fn)
        else:
            fn_for_nat = fn
        yield FunctionInfo.from_fn(
            fn=fn_for_nat,
            description=description,
        )

    return fn