AI Agents templates for trace decorator

Author: antonpirker

sentry_sdk/consts.py

@@ -100,6 +100,17 @@ class CompressionAlgo(Enum):
 ]
 
 
+class SPANTEMPLATE(str, Enum):
+    SPAN = "span"
+    AI_AGENT = "ai_agent"
+    AI_TOOL = "ai_tool"
+    AI_CHAT = "ai_chat"
+
+    def __str__(self):
+        # type: () -> str
+        return self.value
+
+
 class INSTRUMENTER:
     SENTRY = "sentry"
     OTEL = "otel"

sentry_sdk/tracing.py

@@ -5,7 +5,7 @@
 from enum import Enum
 
 import sentry_sdk
-from sentry_sdk.consts import INSTRUMENTER, SPANSTATUS, SPANDATA
+from sentry_sdk.consts import INSTRUMENTER, SPANSTATUS, SPANDATA, SPANTEMPLATE
 from sentry_sdk.profiler.continuous_profiler import get_profiler_id
 from sentry_sdk.utils import (
     get_current_thread_meta,
@@ -1340,8 +1340,10 @@ def _set_initial_sampling_decision(self, sampling_context):
 if TYPE_CHECKING:
 
     @overload
-    def trace(func=None, *, op=None, name=None, attributes=None):
-        # type: (None, Optional[str], Optional[str], Optional[dict[str, Any]]) -> Callable[[Callable[P, R]], Callable[P, R]]
+    def trace(
+        func=None, *, op=None, name=None, attributes=None, template=SPANTEMPLATE.SPAN
+    ):
+        # type: (None, Optional[str], Optional[str], Optional[dict[str, Any]], SPANTEMPLATE) -> Callable[[Callable[P, R]], Callable[P, R]]
         # Handles: @trace() and @trace(op="custom")
         pass
 
@@ -1352,8 +1354,10 @@ def trace(func):
         pass
 
 
-def trace(func=None, *, op=None, name=None, attributes=None):
-    # type: (Optional[Callable[P, R]], Optional[str], Optional[str], Optional[dict[str, Any]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
+def trace(
+    func=None, *, op=None, name=None, attributes=None, template=SPANTEMPLATE.SPAN
+):
+    # type: (Optional[Callable[P, R]], Optional[str], Optional[str], Optional[dict[str, Any]], SPANTEMPLATE) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
     """
     Decorator to start a child span around a function call.
 
@@ -1382,14 +1386,21 @@ def trace(func=None, *, op=None, name=None, attributes=None):
         attributes provide additional context about the span's execution.
     :type attributes: dict[str, Any] or None
 
+    :param template: The type of span to create. This determines what kind of
+        span instrumentation and data collection will be applied. Use predefined
+        constants from :py:class:`sentry_sdk.consts.SPANTEMPLATE`.
+        The default is `SPANTEMPLATE.SPAN` which is the right choice for most
+        use cases.
+    :type template: :py:class:`sentry_sdk.consts.SPANTEMPLATE`
+
     :returns: When used as ``@trace``, returns the decorated function. When used as
         ``@trace(...)`` with parameters, returns a decorator function.
     :rtype: Callable or decorator function
 
     Example::
 
         import sentry_sdk
-        from sentry_sdk.consts import OP
+        from sentry_sdk.consts import OP, SPANTEMPLATE
 
         # Simple usage with default values
         @sentry_sdk.trace
@@ -1406,13 +1417,20 @@ def process_data():
         def make_db_query(sql):
             # Function implementation
             pass
+
+        # With a custom template
+        @sentry_sdk.trace(template=SPANTEMPLATE.AI_TOOL)
+        def calculate_interest_rate(amount, rate, years):
+            # Function implementation
+            pass
     """
     from sentry_sdk.tracing_utils import create_span_decorator
 
     decorator = create_span_decorator(
         op=op,
         name=name,
         attributes=attributes,
+        template=template,
     )
 
     if func: