feat: add agent and tool decorator

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",
 ]

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