sentry-sdk 2.37.1__py2.py3-none-any.whl → 2.39.0__py2.py3-none-any.whl

sentry_sdk/integrations/huggingface_hub.py

@@ -1,24 +1,25 @@
+import inspect
 from functools import wraps
 
-from sentry_sdk import consts
+import sentry_sdk
 from sentry_sdk.ai.monitoring import record_token_usage
 from sentry_sdk.ai.utils import set_data_normalized
-from sentry_sdk.consts import SPANDATA
-
-from typing import Any, Iterable, Callable
-
-import sentry_sdk
-from sentry_sdk.scope import should_send_default_pii
+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.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
 )
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, Iterable
+
 try:
     import huggingface_hub.inference._client
-
-    from huggingface_hub import ChatCompletionStreamOutput, TextGenerationOutput
 except ImportError:
     raise DidNotEnable("Huggingface not installed")
 
@@ -34,15 +35,26 @@ class HuggingfaceHubIntegration(Integration):
     @staticmethod
     def setup_once():
         # type: () -> None
+
+        # Other tasks that can be called: https://huggingface.co/docs/huggingface_hub/guides/inference#supported-providers-and-tasks
         huggingface_hub.inference._client.InferenceClient.text_generation = (
-            _wrap_text_generation(
-                huggingface_hub.inference._client.InferenceClient.text_generation
+            _wrap_huggingface_task(
+                huggingface_hub.inference._client.InferenceClient.text_generation,
+                OP.GEN_AI_GENERATE_TEXT,
+            )
+        )
+        huggingface_hub.inference._client.InferenceClient.chat_completion = (
+            _wrap_huggingface_task(
+                huggingface_hub.inference._client.InferenceClient.chat_completion,
+                OP.GEN_AI_CHAT,
             )
         )
 
 
 def _capture_exception(exc):
     # type: (Any) -> None
+    set_span_errored()
+
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -51,34 +63,70 @@ def _capture_exception(exc):
     sentry_sdk.capture_event(event, hint=hint)
 
 
-def _wrap_text_generation(f):
-    # type: (Callable[..., Any]) -> Callable[..., Any]
+def _wrap_huggingface_task(f, op):
+    # type: (Callable[..., Any], str) -> Callable[..., Any]
     @wraps(f)
-    def new_text_generation(*args, **kwargs):
+    def new_huggingface_task(*args, **kwargs):
         # type: (*Any, **Any) -> Any
         integration = sentry_sdk.get_client().get_integration(HuggingfaceHubIntegration)
         if integration is None:
             return f(*args, **kwargs)
 
+        prompt = None
         if "prompt" in kwargs:
             prompt = kwargs["prompt"]
+        elif "messages" in kwargs:
+            prompt = kwargs["messages"]
         elif len(args) >= 2:
-            kwargs["prompt"] = args[1]
-            prompt = kwargs["prompt"]
-            args = (args[0],) + args[2:]
-        else:
-            # invalid call, let it return error
+            if isinstance(args[1], str) or isinstance(args[1], list):
+                prompt = args[1]
+
+        if prompt is None:
+            # invalid call, dont instrument, let it return error
             return f(*args, **kwargs)
 
-        model = kwargs.get("model")
-        streaming = kwargs.get("stream")
+        client = args[0]
+        model = client.model or kwargs.get("model") or ""
+        operation_name = op.split(".")[-1]
 
         span = sentry_sdk.start_span(
-            op=consts.OP.HUGGINGFACE_HUB_CHAT_COMPLETIONS_CREATE,
-            name="Text Generation",
+            op=op,
+            name=f"{operation_name} {model}",
             origin=HuggingfaceHubIntegration.origin,
         )
         span.__enter__()
+
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, operation_name)
+
+        if model:
+            span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model)
+
+        # Input attributes
+        if should_send_default_pii() and integration.include_prompts:
+            set_data_normalized(
+                span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False
+            )
+
+        attribute_mapping = {
+            "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
+            "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
+            "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS,
+            "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
+            "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE,
+            "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P,
+            "top_k": SPANDATA.GEN_AI_REQUEST_TOP_K,
+            "stream": SPANDATA.GEN_AI_RESPONSE_STREAMING,
+        }
+
+        for attribute, span_attribute in attribute_mapping.items():
+            value = kwargs.get(attribute, None)
+            if value is not None:
+                if isinstance(value, (int, float, bool, str)):
+                    span.set_data(span_attribute, value)
+                else:
+                    set_data_normalized(span, span_attribute, value, unpack=False)
+
+        # LLM Execution
         try:
             res = f(*args, **kwargs)
         except Exception as e:
@@ -86,96 +134,245 @@ def _wrap_text_generation(f):
             span.__exit__(None, None, None)
             raise e from None
 
+        # Output attributes
+        finish_reason = None
+        response_model = None
+        response_text_buffer: list[str] = []
+        tokens_used = 0
+        tool_calls = None
+        usage = None
+
         with capture_internal_exceptions():
-            if should_send_default_pii() and integration.include_prompts:
-                set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, prompt)
+            if isinstance(res, str) and res is not None:
+                response_text_buffer.append(res)
 
-            set_data_normalized(span, SPANDATA.AI_MODEL_ID, model)
-            set_data_normalized(span, SPANDATA.AI_STREAMING, streaming)
+            if hasattr(res, "generated_text") and res.generated_text is not None:
+                response_text_buffer.append(res.generated_text)
 
-            if isinstance(res, str):
-                if should_send_default_pii() and integration.include_prompts:
-                    set_data_normalized(
-                        span,
-                        SPANDATA.AI_RESPONSES,
-                        [res],
-                    )
-                span.__exit__(None, None, None)
-                return res
+            if hasattr(res, "model") and res.model is not None:
+                response_model = res.model
+
+            if hasattr(res, "details") and hasattr(res.details, "finish_reason"):
+                finish_reason = res.details.finish_reason
+
+            if (
+                hasattr(res, "details")
+                and hasattr(res.details, "generated_tokens")
+                and res.details.generated_tokens is not None
+            ):
+                tokens_used = res.details.generated_tokens
+
+            if hasattr(res, "usage") and res.usage is not None:
+                usage = res.usage
+
+            if hasattr(res, "choices") and res.choices is not None:
+                for choice in res.choices:
+                    if hasattr(choice, "finish_reason"):
+                        finish_reason = choice.finish_reason
+                    if hasattr(choice, "message") and hasattr(
+                        choice.message, "tool_calls"
+                    ):
+                        tool_calls = choice.message.tool_calls
+                    if (
+                        hasattr(choice, "message")
+                        and hasattr(choice.message, "content")
+                        and choice.message.content is not None
+                    ):
+                        response_text_buffer.append(choice.message.content)
 
-            if isinstance(res, TextGenerationOutput):
-                if should_send_default_pii() and integration.include_prompts:
+            if response_model is not None:
+                span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
+
+            if finish_reason is not None:
+                set_data_normalized(
+                    span,
+                    SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS,
+                    finish_reason,
+                )
+
+            if should_send_default_pii() and integration.include_prompts:
+                if tool_calls is not None and len(tool_calls) > 0:
                     set_data_normalized(
                         span,
-                        SPANDATA.AI_RESPONSES,
-                        [res.generated_text],
+                        SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+                        tool_calls,
+                        unpack=False,
                     )
-                if res.details is not None and res.details.generated_tokens > 0:
-                    record_token_usage(
-                        span,
-                        total_tokens=res.details.generated_tokens,
-                    )
-                span.__exit__(None, None, None)
-                return res
 
-            if not isinstance(res, Iterable):
-                # we only know how to deal with strings and iterables, ignore
-                set_data_normalized(span, "unknown_response", True)
+                if len(response_text_buffer) > 0:
+                    text_response = "".join(response_text_buffer)
+                    if text_response:
+                        set_data_normalized(
+                            span,
+                            SPANDATA.GEN_AI_RESPONSE_TEXT,
+                            text_response,
+                        )
+
+            if usage is not None:
+                record_token_usage(
+                    span,
+                    input_tokens=usage.prompt_tokens,
+                    output_tokens=usage.completion_tokens,
+                    total_tokens=usage.total_tokens,
+                )
+            elif tokens_used > 0:
+                record_token_usage(
+                    span,
+                    total_tokens=tokens_used,
+                )
+
+            # If the response is not a generator (meaning a streaming response)
+            # we are done and can return the response
+            if not inspect.isgenerator(res):
                 span.__exit__(None, None, None)
                 return res
 
             if kwargs.get("details", False):
-                # res is Iterable[TextGenerationStreamOutput]
+                # text-generation stream output
                 def new_details_iterator():
-                    # type: () -> Iterable[ChatCompletionStreamOutput]
+                    # type: () -> Iterable[Any]
+                    finish_reason = None
+                    response_text_buffer: list[str] = []
+                    tokens_used = 0
+
                     with capture_internal_exceptions():
-                        tokens_used = 0
-                        data_buf: list[str] = []
-                        for x in res:
-                            if hasattr(x, "token") and hasattr(x.token, "text"):
-                                data_buf.append(x.token.text)
-                            if hasattr(x, "details") and hasattr(
-                                x.details, "generated_tokens"
+                        for chunk in res:
+                            if (
+                                hasattr(chunk, "token")
+                                and hasattr(chunk.token, "text")
+                                and chunk.token.text is not None
+                            ):
+                                response_text_buffer.append(chunk.token.text)
+
+                            if hasattr(chunk, "details") and hasattr(
+                                chunk.details, "finish_reason"
+                            ):
+                                finish_reason = chunk.details.finish_reason
+
+                            if (
+                                hasattr(chunk, "details")
+                                and hasattr(chunk.details, "generated_tokens")
+                                and chunk.details.generated_tokens is not None
                             ):
-                                tokens_used = x.details.generated_tokens
-                            yield x
-                        if (
-                            len(data_buf) > 0
-                            and should_send_default_pii()
-                            and integration.include_prompts
-                        ):
+                                tokens_used = chunk.details.generated_tokens
+
+                            yield chunk
+
+                        if finish_reason is not None:
                             set_data_normalized(
-                                span, SPANDATA.AI_RESPONSES, "".join(data_buf)
+                                span,
+                                SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS,
+                                finish_reason,
                             )
+
+                        if should_send_default_pii() and integration.include_prompts:
+                            if len(response_text_buffer) > 0:
+                                text_response = "".join(response_text_buffer)
+                                if text_response:
+                                    set_data_normalized(
+                                        span,
+                                        SPANDATA.GEN_AI_RESPONSE_TEXT,
+                                        text_response,
+                                    )
+
                         if tokens_used > 0:
                             record_token_usage(
                                 span,
                                 total_tokens=tokens_used,
                             )
+
                     span.__exit__(None, None, None)
 
                 return new_details_iterator()
-            else:
-                # res is Iterable[str]
 
+            else:
+                # chat-completion stream output
                 def new_iterator():
                     # type: () -> Iterable[str]
-                    data_buf: list[str] = []
+                    finish_reason = None
+                    response_model = None
+                    response_text_buffer: list[str] = []
+                    tool_calls = None
+                    usage = None
+
                     with capture_internal_exceptions():
-                        for s in res:
-                            if isinstance(s, str):
-                                data_buf.append(s)
-                            yield s
-                        if (
-                            len(data_buf) > 0
-                            and should_send_default_pii()
-                            and integration.include_prompts
-                        ):
+                        for chunk in res:
+                            if hasattr(chunk, "model") and chunk.model is not None:
+                                response_model = chunk.model
+
+                            if hasattr(chunk, "usage") and chunk.usage is not None:
+                                usage = chunk.usage
+
+                            if isinstance(chunk, str):
+                                if chunk is not None:
+                                    response_text_buffer.append(chunk)
+
+                            if hasattr(chunk, "choices") and chunk.choices is not None:
+                                for choice in chunk.choices:
+                                    if (
+                                        hasattr(choice, "delta")
+                                        and hasattr(choice.delta, "content")
+                                        and choice.delta.content is not None
+                                    ):
+                                        response_text_buffer.append(
+                                            choice.delta.content
+                                        )
+
+                                    if (
+                                        hasattr(choice, "finish_reason")
+                                        and choice.finish_reason is not None
+                                    ):
+                                        finish_reason = choice.finish_reason
+
+                                    if (
+                                        hasattr(choice, "delta")
+                                        and hasattr(choice.delta, "tool_calls")
+                                        and choice.delta.tool_calls is not None
+                                    ):
+                                        tool_calls = choice.delta.tool_calls
+
+                            yield chunk
+
+                        if response_model is not None:
+                            span.set_data(
+                                SPANDATA.GEN_AI_RESPONSE_MODEL, response_model
+                            )
+
+                        if finish_reason is not None:
                             set_data_normalized(
-                                span, SPANDATA.AI_RESPONSES, "".join(data_buf)
+                                span,
+                                SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS,
+                                finish_reason,
                             )
+
+                        if should_send_default_pii() and integration.include_prompts:
+                            if tool_calls is not None and len(tool_calls) > 0:
+                                set_data_normalized(
+                                    span,
+                                    SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+                                    tool_calls,
+                                    unpack=False,
+                                )
+
+                            if len(response_text_buffer) > 0:
+                                text_response = "".join(response_text_buffer)
+                                if text_response:
+                                    set_data_normalized(
+                                        span,
+                                        SPANDATA.GEN_AI_RESPONSE_TEXT,
+                                        text_response,
+                                    )
+
+                        if usage is not None:
+                            record_token_usage(
+                                span,
+                                input_tokens=usage.prompt_tokens,
+                                output_tokens=usage.completion_tokens,
+                                total_tokens=usage.total_tokens,
+                            )
+
                         span.__exit__(None, None, None)
 
                 return new_iterator()
 
-    return new_text_generation
+    return new_huggingface_task

sentry_sdk/integrations/langchain.py

@@ -4,12 +4,11 @@ from functools import wraps
 
 import sentry_sdk
 from sentry_sdk.ai.monitoring import set_ai_pipeline_name
-from sentry_sdk.ai.utils import set_data_normalized
+from sentry_sdk.ai.utils import set_data_normalized, get_start_span_function
 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.tracing import Span
-from sentry_sdk.tracing_utils import _get_value
+from sentry_sdk.tracing_utils import _get_value, set_span_errored
 from sentry_sdk.utils import logger, capture_internal_exceptions
 
 from typing import TYPE_CHECKING
@@ -26,6 +25,7 @@ if TYPE_CHECKING:
         Union,
     )
     from uuid import UUID
+    from sentry_sdk.tracing import Span
 
 
 try:
@@ -116,7 +116,7 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
 
             span_data = self.span_map[run_id]
             span = span_data.span
-            span.set_status("unknown")
+            set_span_errored(span)
 
             sentry_sdk.capture_exception(error, span.scope)
 
@@ -322,14 +322,15 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                     pass
 
                 try:
-                    tool_calls = getattr(generation.message, "tool_calls", None)
-                    if tool_calls is not None and tool_calls != []:
-                        set_data_normalized(
-                            span,
-                            SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
-                            tool_calls,
-                            unpack=False,
-                        )
+                    if should_send_default_pii() and self.include_prompts:
+                        tool_calls = getattr(generation.message, "tool_calls", None)
+                        if tool_calls is not None and tool_calls != []:
+                            set_data_normalized(
+                                span,
+                                SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+                                tool_calls,
+                                unpack=False,
+                            )
                 except AttributeError:
                     pass
 
@@ -716,8 +717,9 @@ def _wrap_agent_executor_invoke(f):
             return f(self, *args, **kwargs)
 
         agent_name, tools = _get_request_data(self, args, kwargs)
+        start_span_function = get_start_span_function()
 
-        with sentry_sdk.start_span(
+        with start_span_function(
             op=OP.GEN_AI_INVOKE_AGENT,
             name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
             origin=LangchainIntegration.origin,
@@ -767,8 +769,9 @@ def _wrap_agent_executor_stream(f):
             return f(self, *args, **kwargs)
 
         agent_name, tools = _get_request_data(self, args, kwargs)
+        start_span_function = get_start_span_function()
 
-        span = sentry_sdk.start_span(
+        span = start_span_function(
             op=OP.GEN_AI_INVOKE_AGENT,
             name=f"invoke_agent {agent_name}".strip(),
             origin=LangchainIntegration.origin,

sentry_sdk/integrations/openai.py

@@ -7,6 +7,7 @@ from sentry_sdk.ai.utils import set_data_normalized
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
 from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.tracing_utils import set_span_errored
 from sentry_sdk.utils import (
     capture_internal_exceptions,
     event_from_exception,
@@ -83,6 +84,8 @@ def _capture_exception(exc, manual_span_cleanup=True):
     # Close an eventually open span
     # We need to do this by hand because we are not using the start_span context manager
     current_span = sentry_sdk.get_current_span()
+    set_span_errored(current_span)
+
     if manual_span_cleanup and current_span is not None:
         current_span.__exit__(None, None, None)
 
@@ -279,9 +282,9 @@ def _set_output_data(span, response, kwargs, integration, finish_span=True):
 
         def new_iterator():
             # type: () -> Iterator[ChatCompletionChunk]
-            with capture_internal_exceptions():
-                count_tokens_manually = True
-                for x in old_iterator:
+            count_tokens_manually = True
+            for x in old_iterator:
+                with capture_internal_exceptions():
                     # OpenAI chat completion API
                     if hasattr(x, "choices"):
                         choice_index = 0
@@ -312,8 +315,9 @@ def _set_output_data(span, response, kwargs, integration, finish_span=True):
                         )
                         count_tokens_manually = False
 
-                    yield x
+                yield x
 
+            with capture_internal_exceptions():
                 if len(data_buf) > 0:
                     all_responses = ["".join(chunk) for chunk in data_buf]
                     if should_send_default_pii() and integration.include_prompts:
@@ -334,9 +338,9 @@ def _set_output_data(span, response, kwargs, integration, finish_span=True):
 
         async def new_iterator_async():
             # type: () -> AsyncIterator[ChatCompletionChunk]
-            with capture_internal_exceptions():
-                count_tokens_manually = True
-                async for x in old_iterator:
+            count_tokens_manually = True
+            async for x in old_iterator:
+                with capture_internal_exceptions():
                     # OpenAI chat completion API
                     if hasattr(x, "choices"):
                         choice_index = 0
@@ -367,8 +371,9 @@ def _set_output_data(span, response, kwargs, integration, finish_span=True):
                         )
                         count_tokens_manually = False
 
-                    yield x
+                yield x
 
+            with capture_internal_exceptions():
                 if len(data_buf) > 0:
                     all_responses = ["".join(chunk) for chunk in data_buf]
                     if should_send_default_pii() and integration.include_prompts:

sentry_sdk/integrations/openai_agents/patches/agent_run.py

@@ -26,12 +26,12 @@ def _patch_agent_run():
     original_execute_handoffs = agents._run_impl.RunImpl.execute_handoffs
     original_execute_final_output = agents._run_impl.RunImpl.execute_final_output
 
-    def _start_invoke_agent_span(context_wrapper, agent):
-        # type: (agents.RunContextWrapper, agents.Agent) -> None
+    def _start_invoke_agent_span(context_wrapper, agent, kwargs):
+        # type: (agents.RunContextWrapper, agents.Agent, dict[str, Any]) -> None
         """Start an agent invocation span"""
         # Store the agent on the context wrapper so we can access it later
         context_wrapper._sentry_current_agent = agent
-        invoke_agent_span(context_wrapper, agent)
+        invoke_agent_span(context_wrapper, agent, kwargs)
 
     def _end_invoke_agent_span(context_wrapper, agent, output=None):
         # type: (agents.RunContextWrapper, agents.Agent, Optional[Any]) -> None
@@ -72,7 +72,7 @@ def _patch_agent_run():
                 if current_agent and current_agent != agent:
                     _end_invoke_agent_span(context_wrapper, current_agent)
 
-            _start_invoke_agent_span(context_wrapper, agent)
+            _start_invoke_agent_span(context_wrapper, agent, kwargs)
 
         # Call original method with all the correct parameters
         result = await original_run_single_turn(*args, **kwargs)

sentry_sdk/integrations/openai_agents/spans/agent_workflow.py

@@ -1,7 +1,7 @@
 import sentry_sdk
+from sentry_sdk.ai.utils import get_start_span_function
 
 from ..consts import SPAN_ORIGIN
-from ..utils import _get_start_span_function
 
 from typing import TYPE_CHECKING
 
@@ -13,7 +13,7 @@ def agent_workflow_span(agent):
     # type: (agents.Agent) -> sentry_sdk.tracing.Span
 
     # Create a transaction or a span if an transaction is already active
-    span = _get_start_span_function()(
+    span = get_start_span_function()(
         name=f"{agent.name} workflow",
         origin=SPAN_ORIGIN,
     )

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -42,7 +42,7 @@ def update_execute_tool_span(span, agent, tool, result):
     if isinstance(result, str) and result.startswith(
         "An error occurred while running the tool"
     ):
-        span.set_status(SPANSTATUS.INTERNAL_ERROR)
+        span.set_status(SPANSTATUS.ERROR)
 
     if should_send_default_pii():
         span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, result)