diff --git a/alphatrion/tracing/__init__.py b/alphatrion/tracing/__init__.py index 3f98b5a..12a55b1 100644 --- a/alphatrion/tracing/__init__.py +++ b/alphatrion/tracing/__init__.py @@ -1,6 +1,8 @@ -from alphatrion.tracing.tracing import task, workflow +from alphatrion.tracing.tracing import agent, task, tool, workflow __all__ = [ "task", "workflow", + "tool", + "agent", ] diff --git a/alphatrion/tracing/tracing.py b/alphatrion/tracing/tracing.py index 2cb799a..91f9941 100644 --- a/alphatrion/tracing/tracing.py +++ b/alphatrion/tracing/tracing.py @@ -1,7 +1,9 @@ import uuid from opentelemetry.semconv_ai import TraceloopSpanKindValues +from traceloop.sdk.decorators import agent as _agent from traceloop.sdk.decorators import task as _task +from traceloop.sdk.decorators import tool as _tool from traceloop.sdk.decorators import workflow as _workflow @@ -49,3 +51,41 @@ def decorator(func): )(func) return decorator + + +def tool( + version: int | None = None, + method_name: str | None = None, +): + """Tool decorator for tracing. + + Attributes (run_id, team_id, experiment_id) are automatically + added to all spans by ContextAttributesSpanProcessor. + """ + + def decorator(func): + return _tool( + version=version, + method_name=method_name, + )(func) + + return decorator + + +def agent( + version: int | None = None, + method_name: str | None = None, +): + """Agent decorator for tracing. + + Attributes (run_id, team_id, experiment_id) are automatically + added to all spans by ContextAttributesSpanProcessor. + """ + + def decorator(func): + return _agent( + version=version, + method_name=method_name, + )(func) + + return decorator