First version of helpers for better manual instrumentation

Author: antonpirker

sentry_sdk/__init__.py

@@ -50,6 +50,7 @@
     "start_session",
     "end_session",
     "set_transaction_name",
+    "update_current_span",
 ]
 
 # Initialize the debug support after everything is loaded

sentry_sdk/api.py

@@ -6,7 +6,7 @@
 from sentry_sdk._init_implementation import init
 from sentry_sdk.consts import INSTRUMENTER
 from sentry_sdk.scope import Scope, _ScopeManager, new_scope, isolation_scope
-from sentry_sdk.tracing import NoOpSpan, Transaction, trace
+from sentry_sdk.tracing import NoOpSpan, Transaction, trace, new_trace
 from sentry_sdk.crons import monitor
 
 from typing import TYPE_CHECKING
@@ -81,10 +81,12 @@ def overload(x):
     "start_span",
     "start_transaction",
     "trace",
+    "new_trace",
     "monitor",
     "start_session",
     "end_session",
     "set_transaction_name",
+    "update_current_span",
 ]
 
 
@@ -473,3 +475,25 @@ def end_session():
 def set_transaction_name(name, source=None):
     # type: (str, Optional[str]) -> None
     return get_current_scope().set_transaction_name(name, source)
+
+
+def update_current_span(op=None, name=None, attributes=None):
+    # type: (Optional[str], Optional[str], Optional[dict[str, Union[str, int, float, bool]]]) -> None
+    """
+    Update the current span with the given parameters.
+
+    :param op: The operation of the span.
+    :param name: The name of the span.
+    :param attributes: The attributes of the span.
+    """
+    current_span = get_current_span()
+
+    if op is not None:
+        current_span.op = op
+
+    if name is not None:
+        current_span.name = name
+
+    if attributes is not None:
+        for key, value in attributes.items():
+            current_span.set_data(key, value)

sentry_sdk/tracing.py

@@ -1,11 +1,13 @@
 from decimal import Decimal
+import functools
+import inspect
 import uuid
 import warnings
 from datetime import datetime, timedelta, timezone
 from enum import Enum
 
 import sentry_sdk
-from sentry_sdk.consts import INSTRUMENTER, SPANSTATUS, SPANDATA
+from sentry_sdk.consts import INSTRUMENTER, OP, SPANSTATUS, SPANDATA
 from sentry_sdk.profiler.continuous_profiler import get_profiler_id
 from sentry_sdk.utils import (
     get_current_thread_meta,
@@ -1371,6 +1373,135 @@ async def my_async_function():
         return start_child_span_decorator
 
 
+def set_span_attributes(span, attributes):
+    # type: (Span, dict[str, Any]) -> None
+    for key, value in attributes.items():
+        span.set_data(key, value)
+
+
+def set_input_attributes(span, template, args, kwargs):
+    # depending on `template` set some attributes on the span derived from args and kwargs
+    pass
+
+
+def set_output_attributes(span, template, result):
+    # depending on `template` set some attributes on the span derived from result
+    pass
+
+
+def new_trace(func=None, *, as_type="span", name=None):
+    """
+    Decorator to start a child span.
+
+    :param func: The function to trace.
+    :param as_type: The type of span to create.
+    :param name: The name of the span.
+    """
+
+    def span_decorator(f, *a, **kw):
+        @functools.wraps(f)
+        async def async_wrapper(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
+            op = kw.get("op", OP.FUNCTION)
+            span_name = kw.get("name", f.__name__)
+            attributes = kw.get("attributes", {})
+
+            with sentry_sdk.start_span(
+                op=op,
+                name=span_name,
+            ) as span:
+                set_span_attributes(span, attributes)
+                set_input_attributes(span, as_type, args, kwargs)
+
+                result = await f(*args, **kwargs)
+
+                set_output_attributes(span, as_type, result)
+
+                return result
+
+        @functools.wraps(f)
+        def sync_wrapper(*args, **kwargs):
+            # type: (*Any, **Any) -> Any
+            op = kw.get("op", OP.FUNCTION)
+            span_name = kw.get("name", f.__name__)
+            attributes = kw.get("attributes", {})
+
+            with sentry_sdk.start_span(
+                op=op,
+                name=span_name,
+            ) as span:
+                set_span_attributes(span, attributes)
+                set_input_attributes(span, as_type, args, kwargs)
+
+                result = f(*args, **kwargs)
+
+                set_output_attributes(span, as_type, result)
+
+                return result
+
+        if inspect.iscoroutinefunction(f):
+            wrapper = async_wrapper
+        else:
+            wrapper = sync_wrapper
+
+        return wrapper
+
+    def ai_chat_decorator(f):
+        # type: (Optional[Callable[P, R]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
+        attributes = {
+            "gen_ai.operation.name": "chat",
+        }
+
+        return span_decorator(
+            f,
+            op=OP.GEN_AI_CHAT,
+            name="chat [MODEL]",
+            attributes=attributes,
+        )
+
+    def ai_agent_decorator(f):
+        # type: (Optional[Callable[P, R]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
+        span_name = name or f.__name__
+        attributes = {
+            "gen_ai.agent.name": span_name,
+            "gen_ai.operation.name": "invoke_agent",
+        }
+
+        return span_decorator(
+            f,
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name=f"invoke_agent {span_name}",
+            attributes=attributes,
+        )
+
+    def ai_tool_decorator(f):
+        # type: (Optional[Callable[P, R]]) -> Union[Callable[P, R], Callable[[Callable[P, R]], Callable[P, R]]]
+        span_name = name or f.__name__
+        attributes = {
+            "gen_ai.tool.name": span_name,
+            "gen_ai.operation.name": "execute_tool",
+        }
+
+        return span_decorator(
+            f,
+            op=OP.GEN_AI_EXECUTE_TOOL,
+            name=f"execute_tool {span_name}",
+            attributes=attributes,
+        )
+
+    decorator_for_type = {
+        "ai_chat": ai_chat_decorator,
+        "ai_agent": ai_agent_decorator,
+        "ai_tool": ai_tool_decorator,
+    }
+    decorator = decorator_for_type.get(as_type, span_decorator)
+
+    if func:
+        return decorator(func)
+    else:
+        return decorator
+
+
 # Circular imports
 
 from sentry_sdk.tracing_utils import (