Merge branch 'master' into refactor/dns-check

Author: ericapisani

sentry_sdk/ai/monitoring.py

@@ -2,10 +2,12 @@
 import sys
 from functools import wraps
 
+from sentry_sdk.ai.utils import _set_span_data_attribute
 from sentry_sdk.consts import SPANDATA
 import sentry_sdk.utils
 from sentry_sdk import start_span
 from sentry_sdk.tracing import Span
+from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import ContextVar, reraise, capture_internal_exceptions
 
 from typing import TYPE_CHECKING
@@ -97,7 +99,7 @@ async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any":
 
 
 def record_token_usage(
-    span: "Span",
+    span: "Union[Span, StreamedSpan]",
     input_tokens: "Optional[int]" = None,
     input_tokens_cached: "Optional[int]" = None,
     input_tokens_cache_write: "Optional[int]" = None,
@@ -108,28 +110,33 @@ def record_token_usage(
     # TODO: move pipeline name elsewhere
     ai_pipeline_name = get_ai_pipeline_name()
     if ai_pipeline_name:
-        span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, ai_pipeline_name)
+        _set_span_data_attribute(span, SPANDATA.GEN_AI_PIPELINE_NAME, ai_pipeline_name)
 
     if input_tokens is not None:
-        span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
+        _set_span_data_attribute(span, SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
 
     if input_tokens_cached is not None:
-        span.set_data(
+        _set_span_data_attribute(
+            span,
             SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED,
             input_tokens_cached,
         )
 
     if input_tokens_cache_write is not None:
-        span.set_data(
+        _set_span_data_attribute(
+            span,
             SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE,
             input_tokens_cache_write,
         )
 
     if output_tokens is not None:
-        span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
+        _set_span_data_attribute(
+            span, SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens
+        )
 
     if output_tokens_reasoning is not None:
-        span.set_data(
+        _set_span_data_attribute(
+            span,
             SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING,
             output_tokens_reasoning,
         )
@@ -138,4 +145,4 @@ def record_token_usage(
         total_tokens = input_tokens + output_tokens
 
     if total_tokens is not None:
-        span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)
+        _set_span_data_attribute(span, SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)

sentry_sdk/ai/utils.py

@@ -7,7 +7,7 @@
 from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Dict, List, Optional, Tuple
+    from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
     from sentry_sdk.tracing import Span
 
@@ -490,13 +490,25 @@ def _normalize_data(data: "Any", unpack: bool = True) -> "Any":
 
 
 def set_data_normalized(
-    span: "Span", key: str, value: "Any", unpack: bool = True
+    span: "Union[Span, StreamedSpan]",
+    key: str,
+    value: "Any",
+    unpack: bool = True,
 ) -> None:
     normalized = _normalize_data(value, unpack=unpack)
     if isinstance(normalized, (int, float, bool, str)):
-        span.set_data(key, normalized)
+        _set_span_data_attribute(span, key, normalized)
     else:
-        span.set_data(key, json.dumps(normalized))
+        _set_span_data_attribute(span, key, json.dumps(normalized))
+
+
+def _set_span_data_attribute(
+    span: "Union[Span, StreamedSpan]", key: str, value: "Any"
+) -> None:
+    if isinstance(span, StreamedSpan):
+        span.set_attribute(key, value)
+    else:
+        span.set_data(key, value)
 
 
 def normalize_message_role(role: str) -> str:

sentry_sdk/consts.py

@@ -549,6 +549,12 @@ class SPANDATA:
     Example: "Hello!"
     """
 
+    GEN_AI_FUNCTION_ID = "gen_ai.function_id"
+    """
+    Framework-specific tracing label for the execution of a function or other unit of execution in a generative AI system.
+    Example: "my-awesome-function"
+    """
+
     GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
     """
     The name of the operation being performed.

sentry_sdk/integrations/huggingface_hub.py

@@ -5,18 +5,22 @@
 
 import sentry_sdk
 from sentry_sdk.ai.monitoring import record_token_usage
-from sentry_sdk.ai.utils import set_data_normalized
+from sentry_sdk.ai.utils import _set_span_data_attribute, set_data_normalized
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
     reraise,
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Iterable
+    from typing import Any, Callable, Iterable, Union
+
+    from sentry_sdk.tracing import Span
 
 try:
     import huggingface_hub.inference._client
@@ -83,17 +87,27 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
         model = client.model or kwargs.get("model") or ""
         operation_name = op.split(".")[-1]
 
-        span = sentry_sdk.start_span(
-            op=op,
-            name=f"{operation_name} {model}",
-            origin=HuggingfaceHubIntegration.origin,
-        )
+        span: "Union[Span, StreamedSpan]"
+        if has_span_streaming_enabled(sentry_sdk.get_client().options):
+            span = sentry_sdk.traces.start_span(
+                name=f"{operation_name} {model}",
+                attributes={
+                    "sentry.op": op,
+                    "sentry.origin": HuggingfaceHubIntegration.origin,
+                },
+            )
+        else:
+            span = sentry_sdk.start_span(
+                op=op,
+                name=f"{operation_name} {model}",
+                origin=HuggingfaceHubIntegration.origin,
+            )
         span.__enter__()
 
-        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, operation_name)
+        _set_span_data_attribute(span, SPANDATA.GEN_AI_OPERATION_NAME, operation_name)
 
         if model:
-            span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model)
+            _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model)
 
         # Input attributes
         if should_send_default_pii() and integration.include_prompts:
@@ -116,7 +130,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
             value = kwargs.get(attribute, None)
             if value is not None:
                 if isinstance(value, (int, float, bool, str)):
-                    span.set_data(span_attribute, value)
+                    _set_span_data_attribute(span, span_attribute, value)
                 else:
                     set_data_normalized(span, span_attribute, value, unpack=False)
 
@@ -177,7 +191,9 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
                         response_text_buffer.append(choice.message.content)
 
             if response_model is not None:
-                span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
+                _set_span_data_attribute(
+                    span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model
+                )
 
             if finish_reason is not None:
                 set_data_normalized(
@@ -328,8 +344,8 @@ def new_iterator() -> "Iterable[str]":
                             yield chunk
 
                         if response_model is not None:
-                            span.set_data(
-                                SPANDATA.GEN_AI_RESPONSE_MODEL, response_model
+                            _set_span_data_attribute(
+                                span, SPANDATA.GEN_AI_RESPONSE_MODEL, response_model
                             )
 
                         if finish_reason is not None: