Source code for taimoe.platform.decorators

"""Function-tracking decorators (sync **and** async aware).

``track_action`` wraps an individual tool call inside the current trace;
no trace is auto-created. ``track_agent`` wraps an agent entry point and
will mint a trace if the caller doesn't already have one open.

Both flavors detect ``async def`` callables via ``inspect`` and produce a
correctly-shaped wrapper, so callers don't need separate ``track_action_async``
or ``track_agent_async`` symbols.
"""

from __future__ import annotations

from typing import Any, Callable, TypeVar

from .instrumentation._helpers import with_span

F = TypeVar("F", bound=Callable[..., Any])




def track_action(name: str | None = None) -> Callable[[F], F]:
    """Wrap a function as a ``tool`` span on the active trace.

    If there is no active trace the call passes through unwrapped (we
    don't silently mint a trace for a single action — that would obscure
    which agent owns the action). Works for sync and ``async def`` callables.
    """

    def decorator(func: F) -> F:
        action_name = name or func.__name__
        return with_span(
            span_type="tool",
            name=action_name,
            create_trace_if_missing=False,
        )(func)

    return decorator





def track_agent(name: str | None = None) -> Callable[[F], F]:
    """Wrap an agent entry point as an ``agent`` span, minting a trace if
    needed.

    Works for sync and ``async def`` callables.
    """

    def decorator(func: F) -> F:
        agent_name = name or func.__name__
        return with_span(
            span_type="agent",
            name=agent_name,
            create_trace_if_missing=True,
            trace_name=f"agent-run-{agent_name}",
            agent_id=agent_name,
        )(func)

    return decorator