sentry-sdk 3.0.0a5__py2.py3-none-any.whl → 3.0.0a6__py2.py3-none-any.whl

sentry_sdk/tracing_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import contextlib
+import functools
 import inspect
 import os
 import re
@@ -8,7 +9,6 @@ import uuid
 from collections.abc import Mapping
 from datetime import datetime, timedelta, timezone
 from decimal import ROUND_DOWN, Decimal, DefaultContext, localcontext
-from functools import wraps
 from random import Random
 from urllib.parse import quote, unquote
 
@@ -17,6 +17,7 @@ from sentry_sdk.consts import (
     OP,
     SPANDATA,
     SPANSTATUS,
+    SPANTEMPLATE,
     BAGGAGE_HEADER_NAME,
     SENTRY_TRACE_HEADER_NAME,
 )
@@ -27,6 +28,7 @@ from sentry_sdk.utils import (
     logger,
     match_regex_list,
     qualname_from_function,
+    safe_repr,
     to_string,
     is_sentry_url,
     _is_external_source,
@@ -686,71 +688,117 @@ def normalize_incoming_data(incoming_data: Dict[str, Any]) -> Dict[str, Any]:
     return data
 
 
-def start_child_span_decorator(func: Any) -> Any:
+def create_span_decorator(
+    op: Optional[Union[str, OP]] = None,
+    name: Optional[str] = None,
+    attributes: Optional[dict[str, Any]] = None,
+    template: SPANTEMPLATE = SPANTEMPLATE.DEFAULT,
+) -> Any:
     """
-    Decorator to add child spans for functions.
-
-    See also ``sentry_sdk.tracing.trace()``.
+    Create a span decorator that can wrap both sync and async functions.
+
+    :param op: The operation type for the span.
+    :type op: str or :py:class:`sentry_sdk.consts.OP` or None
+    :param name: The name of the span.
+    :type name: str or None
+    :param attributes: Additional attributes to set on the span.
+    :type attributes: dict 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.DEFAULT` which is the right choice for most
+        use cases.
+    :type template: :py:class:`sentry_sdk.consts.SPANTEMPLATE`
     """
-    # Asynchronous case
-    if inspect.iscoroutinefunction(func):
+    from sentry_sdk.scope import should_send_default_pii
+
+    def span_decorator(f: Any) -> Any:
+        """
+        Decorator to create a span for the given function.
+        """
 
-        @wraps(func)
-        async def func_with_tracing(*args: Any, **kwargs: Any) -> Any:
+        @functools.wraps(f)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
             span = get_current_span()
 
             if span is None:
                 logger.debug(
                     "Cannot create a child span for %s. "
                     "Please start a Sentry transaction before calling this function.",
-                    qualname_from_function(func),
+                    qualname_from_function(f),
                 )
-                return await func(*args, **kwargs)
+                return await f(*args, **kwargs)
+
+            span_op = op or _get_span_op(template)
+            function_name = name or qualname_from_function(f) or ""
+            span_name = _get_span_name(template, function_name, kwargs)
+            send_pii = should_send_default_pii()
 
             with sentry_sdk.start_span(
-                op=OP.FUNCTION,
-                name=qualname_from_function(func),
+                op=span_op,
+                name=span_name,
                 only_as_child_span=True,
-            ):
-                return await func(*args, **kwargs)
+            ) as span:
+                span.set_attributes(attributes or {})
+                _set_input_attributes(
+                    span, template, send_pii, function_name, f, args, kwargs
+                )
+
+                result = await f(*args, **kwargs)
+
+                _set_output_attributes(span, template, send_pii, result)
+
+                return result
 
         try:
-            func_with_tracing.__signature__ = inspect.signature(  # type: ignore[attr-defined]
-                func
-            )
+            async_wrapper.__signature__ = inspect.signature(f)  # type: ignore[attr-defined]
         except Exception:
             pass
 
-    # Synchronous case
-    else:
-
-        @wraps(func)
-        def func_with_tracing(*args: Any, **kwargs: Any) -> Any:
+        @functools.wraps(f)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
             span = get_current_span()
 
             if span is None:
                 logger.debug(
                     "Cannot create a child span for %s. "
                     "Please start a Sentry transaction before calling this function.",
-                    qualname_from_function(func),
+                    qualname_from_function(f),
                 )
-                return func(*args, **kwargs)
+                return f(*args, **kwargs)
+
+            span_op = op or _get_span_op(template)
+            function_name = name or qualname_from_function(f) or ""
+            span_name = _get_span_name(template, function_name, kwargs)
+            send_pii = should_send_default_pii()
 
             with sentry_sdk.start_span(
-                op=OP.FUNCTION,
-                name=qualname_from_function(func),
+                op=span_op,
+                name=span_name,
                 only_as_child_span=True,
-            ):
-                return func(*args, **kwargs)
+            ) as span:
+                span.set_attributes(attributes or {})
+                _set_input_attributes(
+                    span, template, send_pii, function_name, f, args, kwargs
+                )
+
+                result = f(*args, **kwargs)
+
+                _set_output_attributes(span, template, send_pii, result)
+
+                return result
 
         try:
-            func_with_tracing.__signature__ = inspect.signature(  # type: ignore[attr-defined]
-                func
-            )
+            sync_wrapper.__signature__ = inspect.signature(f)  # type: ignore[attr-defined]
         except Exception:
             pass
 
-    return func_with_tracing
+        if inspect.iscoroutinefunction(f):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    return span_decorator
 
 
 def get_current_span(
@@ -813,6 +861,255 @@ def _sample_rand_range(
         return sample_rate, 1.0
 
 
+def _get_value(source: Any, key: str) -> Optional[Any]:
+    """
+    Gets a value from a source object. The source can be a dict or an object.
+    It is checked for dictionary keys and object attributes.
+    """
+    value = None
+    if isinstance(source, dict):
+        value = source.get(key)
+    else:
+        if hasattr(source, key):
+            try:
+                value = getattr(source, key)
+            except Exception:
+                value = None
+    return value
+
+
+def _get_span_name(
+    template: Union[str, SPANTEMPLATE],
+    name: str,
+    kwargs: Optional[dict[str, Any]] = None,
+) -> str:
+    """
+    Get the name of the span based on the template and the name.
+    """
+    span_name = name
+
+    if template == SPANTEMPLATE.AI_CHAT:
+        model = None
+        if kwargs:
+            for key in ("model", "model_name"):
+                if kwargs.get(key) and isinstance(kwargs[key], str):
+                    model = kwargs[key]
+                    break
+
+        span_name = f"chat {model}" if model else "chat"
+
+    elif template == SPANTEMPLATE.AI_AGENT:
+        span_name = f"invoke_agent {name}"
+
+    elif template == SPANTEMPLATE.AI_TOOL:
+        span_name = f"execute_tool {name}"
+
+    return span_name
+
+
+def _get_span_op(template: Union[str, SPANTEMPLATE]) -> str:
+    """
+    Get the operation of the span based on the template.
+    """
+    mapping = {
+        SPANTEMPLATE.AI_CHAT: OP.GEN_AI_CHAT,
+        SPANTEMPLATE.AI_AGENT: OP.GEN_AI_INVOKE_AGENT,
+        SPANTEMPLATE.AI_TOOL: OP.GEN_AI_EXECUTE_TOOL,
+    }  # type: dict[Union[str, SPANTEMPLATE], Union[str, OP]]
+    op = mapping.get(template, OP.FUNCTION)
+
+    return str(op)
+
+
+def _get_input_attributes(
+    template: Union[str, SPANTEMPLATE],
+    send_pii: bool,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    """
+    Get input attributes for the given span template.
+    """
+    attributes: dict[str, Any] = {}
+
+    if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]:
+        mapping: dict[str, tuple[str, type]] = {
+            "model": (SPANDATA.GEN_AI_REQUEST_MODEL, str),
+            "model_name": (SPANDATA.GEN_AI_REQUEST_MODEL, str),
+            "agent": (SPANDATA.GEN_AI_AGENT_NAME, str),
+            "agent_name": (SPANDATA.GEN_AI_AGENT_NAME, str),
+            "max_tokens": (SPANDATA.GEN_AI_REQUEST_MAX_TOKENS, int),
+            "frequency_penalty": (SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY, float),
+            "presence_penalty": (SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY, float),
+            "temperature": (SPANDATA.GEN_AI_REQUEST_TEMPERATURE, float),
+            "top_p": (SPANDATA.GEN_AI_REQUEST_TOP_P, float),
+            "top_k": (SPANDATA.GEN_AI_REQUEST_TOP_K, int),
+        }
+
+        def _set_from_key(key: str, value: Any) -> None:
+            if key in mapping:
+                (attribute, data_type) = mapping[key]
+                if value is not None and isinstance(value, data_type):
+                    attributes[attribute] = value
+
+        for key, value in list(kwargs.items()):
+            if key == "prompt" and isinstance(value, str):
+                attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append(
+                    {"role": "user", "content": value}
+                )
+                continue
+
+            if key == "system_prompt" and isinstance(value, str):
+                attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append(
+                    {"role": "system", "content": value}
+                )
+                continue
+
+            _set_from_key(key, value)
+
+    if template == SPANTEMPLATE.AI_TOOL and send_pii:
+        attributes[SPANDATA.GEN_AI_TOOL_INPUT] = safe_repr(
+            {"args": args, "kwargs": kwargs}
+        )
+
+    # Coerce to string
+    if SPANDATA.GEN_AI_REQUEST_MESSAGES in attributes:
+        attributes[SPANDATA.GEN_AI_REQUEST_MESSAGES] = safe_repr(
+            attributes[SPANDATA.GEN_AI_REQUEST_MESSAGES]
+        )
+
+    return attributes
+
+
+def _get_usage_attributes(usage: Any) -> dict[str, Any]:
+    """
+    Get usage attributes.
+    """
+    attributes = {}
+
+    def _set_from_keys(attribute: str, keys: tuple[str, ...]) -> None:
+        for key in keys:
+            value = _get_value(usage, key)
+            if value is not None and isinstance(value, int):
+                attributes[attribute] = value
+
+    _set_from_keys(
+        SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,
+        ("prompt_tokens", "input_tokens"),
+    )
+    _set_from_keys(
+        SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS,
+        ("completion_tokens", "output_tokens"),
+    )
+    _set_from_keys(
+        SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
+        ("total_tokens",),
+    )
+
+    return attributes
+
+
+def _get_output_attributes(
+    template: Union[str, SPANTEMPLATE], send_pii: bool, result: Any
+) -> dict[str, Any]:
+    """
+    Get output attributes for the given span template.
+    """
+    attributes = {}  # type: dict[str, Any]
+
+    if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]:
+        with capture_internal_exceptions():
+            # Usage from result, result.usage, and result.metadata.usage
+            usage_candidates = [result]
+
+            usage = _get_value(result, "usage")
+            usage_candidates.append(usage)
+
+            meta = _get_value(result, "metadata")
+            usage = _get_value(meta, "usage")
+            usage_candidates.append(usage)
+
+            for usage_candidate in usage_candidates:
+                if usage_candidate is not None:
+                    attributes.update(_get_usage_attributes(usage_candidate))
+
+            # Response model
+            model_name = _get_value(result, "model")
+            if model_name is not None and isinstance(model_name, str):
+                attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = model_name
+
+            model_name = _get_value(result, "model_name")
+            if model_name is not None and isinstance(model_name, str):
+                attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = model_name
+
+    # Tool output
+    if template == SPANTEMPLATE.AI_TOOL and send_pii:
+        attributes[SPANDATA.GEN_AI_TOOL_OUTPUT] = safe_repr(result)
+
+    return attributes
+
+
+def _set_input_attributes(
+    span: sentry_sdk.tracing.Span,
+    template: Union[str, SPANTEMPLATE],
+    send_pii: bool,
+    name: str,
+    f: Any,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> None:
+    """
+    Set span input attributes based on the given span template.
+
+    :param span: The span to set attributes on.
+    :param template: The template to use to set attributes on the span.
+    :param send_pii: Whether to send PII data.
+    :param f: The wrapped function.
+    :param args: The arguments to the wrapped function.
+    :param kwargs: The keyword arguments to the wrapped function.
+    """
+    attributes = {}  # type: dict[str, Any]
+
+    if template == SPANTEMPLATE.AI_AGENT:
+        attributes = {
+            SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent",
+            SPANDATA.GEN_AI_AGENT_NAME: name,
+        }
+    elif template == SPANTEMPLATE.AI_CHAT:
+        attributes = {
+            SPANDATA.GEN_AI_OPERATION_NAME: "chat",
+        }
+    elif template == SPANTEMPLATE.AI_TOOL:
+        attributes = {
+            SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool",
+            SPANDATA.GEN_AI_TOOL_NAME: name,
+        }
+
+        docstring = f.__doc__
+        if docstring is not None:
+            attributes[SPANDATA.GEN_AI_TOOL_DESCRIPTION] = docstring
+
+    attributes.update(_get_input_attributes(template, send_pii, args, kwargs))
+    span.set_attributes(attributes or {})
+
+
+def _set_output_attributes(
+    span: sentry_sdk.tracing.Span,
+    template: Union[str, SPANTEMPLATE],
+    send_pii: bool,
+    result: Any,
+) -> None:
+    """
+    Set span output attributes based on the given span template.
+
+    :param span: The span to set attributes on.
+    :param template: The template to use to set attributes on the span.
+    :param send_pii: Whether to send PII data.
+    :param result: The result of the wrapped function.
+    """
+    span.set_attributes(_get_output_attributes(template, send_pii, result) or {})
+
+
 def get_span_status_from_http_code(http_status_code: int) -> str:
     """
     Returns the Sentry status corresponding to the given HTTP status code.