sentry-sdk 3.0.0a2__py2.py3-none-any.whl → 3.0.0a3__py2.py3-none-any.whl

sentry_sdk/integrations/huey.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import sys
 from datetime import datetime
 
@@ -45,19 +46,16 @@ class HueyIntegration(Integration):
     origin = f"auto.queue.{identifier}"
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         patch_enqueue()
         patch_execute()
 
 
-def patch_enqueue():
-    # type: () -> None
+def patch_enqueue() -> None:
     old_enqueue = Huey.enqueue
 
     @ensure_integration_enabled(HueyIntegration, old_enqueue)
-    def _sentry_enqueue(self, task):
-        # type: (Huey, Task) -> Optional[Union[Result, ResultGroup]]
+    def _sentry_enqueue(self: Huey, task: Task) -> Optional[Union[Result, ResultGroup]]:
         with sentry_sdk.start_span(
             op=OP.QUEUE_SUBMIT_HUEY,
             name=task.name,
@@ -77,10 +75,8 @@ def patch_enqueue():
     Huey.enqueue = _sentry_enqueue
 
 
-def _make_event_processor(task):
-    # type: (Any) -> EventProcessor
-    def event_processor(event, hint):
-        # type: (Event, Hint) -> Optional[Event]
+def _make_event_processor(task: Any) -> EventProcessor:
+    def event_processor(event: Event, hint: Hint) -> Optional[Event]:
 
         with capture_internal_exceptions():
             tags = event.setdefault("tags", {})
@@ -107,8 +103,7 @@ def _make_event_processor(task):
     return event_processor
 
 
-def _capture_exception(exc_info):
-    # type: (ExcInfo) -> None
+def _capture_exception(exc_info: ExcInfo) -> None:
     scope = sentry_sdk.get_current_scope()
 
     if scope.root_span is not None:
@@ -126,12 +121,10 @@ def _capture_exception(exc_info):
     scope.capture_event(event, hint=hint)
 
 
-def _wrap_task_execute(func):
-    # type: (F) -> F
+def _wrap_task_execute(func: F) -> F:
 
     @ensure_integration_enabled(HueyIntegration, func)
-    def _sentry_execute(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    def _sentry_execute(*args: Any, **kwargs: Any) -> Any:
         try:
             result = func(*args, **kwargs)
         except Exception:
@@ -148,13 +141,13 @@ def _wrap_task_execute(func):
     return _sentry_execute  # type: ignore
 
 
-def patch_execute():
-    # type: () -> None
+def patch_execute() -> None:
     old_execute = Huey._execute
 
     @ensure_integration_enabled(HueyIntegration, old_execute)
-    def _sentry_execute(self, task, timestamp=None):
-        # type: (Huey, Task, Optional[datetime]) -> Any
+    def _sentry_execute(
+        self: Huey, task: Task, timestamp: Optional[datetime] = None
+    ) -> Any:
         with sentry_sdk.isolation_scope() as scope:
             with capture_internal_exceptions():
                 scope._name = "huey"

sentry_sdk/integrations/huggingface_hub.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from functools import wraps
 
 from sentry_sdk import consts
@@ -27,13 +28,11 @@ class HuggingfaceHubIntegration(Integration):
     identifier = "huggingface_hub"
     origin = f"auto.ai.{identifier}"
 
-    def __init__(self, include_prompts=True):
-        # type: (HuggingfaceHubIntegration, bool) -> None
+    def __init__(self: HuggingfaceHubIntegration, include_prompts: bool = True) -> None:
         self.include_prompts = include_prompts
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         huggingface_hub.inference._client.InferenceClient.text_generation = (
             _wrap_text_generation(
                 huggingface_hub.inference._client.InferenceClient.text_generation
@@ -41,8 +40,7 @@ class HuggingfaceHubIntegration(Integration):
         )
 
 
-def _capture_exception(exc):
-    # type: (Any) -> None
+def _capture_exception(exc: Any) -> None:
     event, hint = event_from_exception(
         exc,
         client_options=sentry_sdk.get_client().options,
@@ -51,11 +49,9 @@ def _capture_exception(exc):
     sentry_sdk.capture_event(event, hint=hint)
 
 
-def _wrap_text_generation(f):
-    # type: (Callable[..., Any]) -> Callable[..., Any]
+def _wrap_text_generation(f: Callable[..., Any]) -> Callable[..., Any]:
     @wraps(f)
-    def new_text_generation(*args, **kwargs):
-        # type: (*Any, **Any) -> Any
+    def new_text_generation(*args: Any, **kwargs: Any) -> Any:
         integration = sentry_sdk.get_client().get_integration(HuggingfaceHubIntegration)
         if integration is None:
             return f(*args, **kwargs)
@@ -112,7 +108,10 @@ def _wrap_text_generation(f):
                         [res.generated_text],
                     )
                 if res.details is not None and res.details.generated_tokens > 0:
-                    record_token_usage(span, total_tokens=res.details.generated_tokens)
+                    record_token_usage(
+                        span,
+                        total_tokens=res.details.generated_tokens,
+                    )
                 span.__exit__(None, None, None)
                 return res
 
@@ -124,8 +123,7 @@ def _wrap_text_generation(f):
 
             if kwargs.get("details", False):
                 # res is Iterable[TextGenerationStreamOutput]
-                def new_details_iterator():
-                    # type: () -> Iterable[ChatCompletionStreamOutput]
+                def new_details_iterator() -> Iterable[ChatCompletionStreamOutput]:
                     with capture_internal_exceptions():
                         tokens_used = 0
                         data_buf: list[str] = []
@@ -146,15 +144,17 @@ def _wrap_text_generation(f):
                                 span, SPANDATA.AI_RESPONSES, "".join(data_buf)
                             )
                         if tokens_used > 0:
-                            record_token_usage(span, total_tokens=tokens_used)
+                            record_token_usage(
+                                span,
+                                total_tokens=tokens_used,
+                            )
                     span.__exit__(None, None, None)
 
                 return new_details_iterator()
             else:
                 # res is Iterable[str]
 
-                def new_iterator():
-                    # type: () -> Iterable[str]
+                def new_iterator() -> Iterable[str]:
                     data_buf: list[str] = []
                     with capture_internal_exceptions():
                         for s in res:

sentry_sdk/integrations/langchain.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+import itertools
 from collections import OrderedDict
 from functools import wraps
 
@@ -22,6 +24,8 @@ try:
     from langchain_core.callbacks import (
         manager,
         BaseCallbackHandler,
+        BaseCallbackManager,
+        Callbacks,
     )
     from langchain_core.agents import AgentAction, AgentFinish
 except ImportError:
@@ -58,40 +62,41 @@ class LangchainIntegration(Integration):
     max_spans = 1024
 
     def __init__(
-        self, include_prompts=True, max_spans=1024, tiktoken_encoding_name=None
-    ):
-        # type: (LangchainIntegration, bool, int, Optional[str]) -> None
+        self: LangchainIntegration,
+        include_prompts: bool = True,
+        max_spans: int = 1024,
+        tiktoken_encoding_name: Optional[str] = None,
+    ) -> None:
         self.include_prompts = include_prompts
         self.max_spans = max_spans
         self.tiktoken_encoding_name = tiktoken_encoding_name
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         manager._configure = _wrap_configure(manager._configure)
 
 
 class WatchedSpan:
-    num_completion_tokens = 0  # type: int
-    num_prompt_tokens = 0  # type: int
-    no_collect_tokens = False  # type: bool
-    children = []  # type: List[WatchedSpan]
-    is_pipeline = False  # type: bool
-
-    def __init__(self, span):
-        # type: (Span) -> None
+    num_completion_tokens: int = 0
+    num_prompt_tokens: int = 0
+    no_collect_tokens: bool = False
+    children: List[WatchedSpan] = []
+    is_pipeline: bool = False
+
+    def __init__(self, span: Span) -> None:
         self.span = span
 
 
 class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
     """Base callback handler that can be used to handle callbacks from langchain."""
 
-    span_map = OrderedDict()  # type: OrderedDict[UUID, WatchedSpan]
-
-    max_span_map_size = 0
-
-    def __init__(self, max_span_map_size, include_prompts, tiktoken_encoding_name=None):
-        # type: (int, bool, Optional[str]) -> None
+    def __init__(
+        self,
+        max_span_map_size: int,
+        include_prompts: bool,
+        tiktoken_encoding_name: Optional[str] = None,
+    ) -> None:
+        self.span_map: OrderedDict[UUID, WatchedSpan] = OrderedDict()
         self.max_span_map_size = max_span_map_size
         self.include_prompts = include_prompts
 
@@ -101,21 +106,18 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
 
             self.tiktoken_encoding = tiktoken.get_encoding(tiktoken_encoding_name)
 
-    def count_tokens(self, s):
-        # type: (str) -> int
+    def count_tokens(self, s: str) -> int:
         if self.tiktoken_encoding is not None:
             return len(self.tiktoken_encoding.encode_ordinary(s))
         return 0
 
-    def gc_span_map(self):
-        # type: () -> None
+    def gc_span_map(self) -> None:
 
         while len(self.span_map) > self.max_span_map_size:
             run_id, watched_span = self.span_map.popitem(last=False)
             self._exit_span(watched_span, run_id)
 
-    def _handle_error(self, run_id, error):
-        # type: (UUID, Any) -> None
+    def _handle_error(self, run_id: UUID, error: Any) -> None:
         if not run_id or run_id not in self.span_map:
             return
 
@@ -127,14 +129,17 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
         span_data.span.finish()
         del self.span_map[run_id]
 
-    def _normalize_langchain_message(self, message):
-        # type: (BaseMessage) -> Any
+    def _normalize_langchain_message(self, message: BaseMessage) -> Any:
         parsed = {"content": message.content, "role": message.type}
         parsed.update(message.additional_kwargs)
         return parsed
 
-    def _create_span(self, run_id, parent_id, **kwargs):
-        # type: (SentryLangchainCallback, UUID, Optional[Any], Any) -> WatchedSpan
+    def _create_span(
+        self: SentryLangchainCallback,
+        run_id: UUID,
+        parent_id: Optional[Any],
+        **kwargs: Any,
+    ) -> WatchedSpan:
 
         parent_watched_span = self.span_map.get(parent_id) if parent_id else None
         sentry_span = sentry_sdk.start_span(
@@ -161,8 +166,9 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
         self.gc_span_map()
         return watched_span
 
-    def _exit_span(self, span_data, run_id):
-        # type: (SentryLangchainCallback, WatchedSpan, UUID) -> None
+    def _exit_span(
+        self: SentryLangchainCallback, span_data: WatchedSpan, run_id: UUID
+    ) -> None:
 
         if span_data.is_pipeline:
             set_ai_pipeline_name(None)
@@ -172,17 +178,16 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
         del self.span_map[run_id]
 
     def on_llm_start(
-        self,
-        serialized,
-        prompts,
+        self: SentryLangchainCallback,
+        serialized: Dict[str, Any],
+        prompts: List[str],
         *,
-        run_id,
-        tags=None,
-        parent_run_id=None,
-        metadata=None,
-        **kwargs,
-    ):
-        # type: (SentryLangchainCallback, Dict[str, Any], List[str], UUID, Optional[List[str]], Optional[UUID], Optional[Dict[str, Any]], Any) -> Any
+        run_id: UUID,
+        tags: Optional[List[str]] = None,
+        parent_run_id: Optional[UUID] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
         """Run when LLM starts running."""
         with capture_internal_exceptions():
             if not run_id:
@@ -203,8 +208,14 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                 if k in all_params:
                     set_data_normalized(span, v, all_params[k])
 
-    def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, Dict[str, Any], List[List[BaseMessage]], UUID, Any) -> Any
+    def on_chat_model_start(
+        self: SentryLangchainCallback,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when Chat Model starts running."""
         with capture_internal_exceptions():
             if not run_id:
@@ -249,8 +260,9 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                             message.content
                         ) + self.count_tokens(message.type)
 
-    def on_llm_new_token(self, token, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, str, UUID, Any) -> Any
+    def on_llm_new_token(
+        self: SentryLangchainCallback, token: str, *, run_id: UUID, **kwargs: Any
+    ) -> Any:
         """Run on new LLM token. Only available when streaming is enabled."""
         with capture_internal_exceptions():
             if not run_id or run_id not in self.span_map:
@@ -260,8 +272,13 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                 return
             span_data.num_completion_tokens += self.count_tokens(token)
 
-    def on_llm_end(self, response, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
+    def on_llm_end(
+        self: SentryLangchainCallback,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when LLM ends running."""
         with capture_internal_exceptions():
             if not run_id:
@@ -286,27 +303,38 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                 if token_usage:
                     record_token_usage(
                         span_data.span,
-                        token_usage.get("prompt_tokens"),
-                        token_usage.get("completion_tokens"),
-                        token_usage.get("total_tokens"),
+                        input_tokens=token_usage.get("prompt_tokens"),
+                        output_tokens=token_usage.get("completion_tokens"),
+                        total_tokens=token_usage.get("total_tokens"),
                     )
                 else:
                     record_token_usage(
                         span_data.span,
-                        span_data.num_prompt_tokens,
-                        span_data.num_completion_tokens,
+                        input_tokens=span_data.num_prompt_tokens,
+                        output_tokens=span_data.num_completion_tokens,
                     )
 
             self._exit_span(span_data, run_id)
 
-    def on_llm_error(self, error, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
+    def on_llm_error(
+        self: SentryLangchainCallback,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when LLM errors."""
         with capture_internal_exceptions():
             self._handle_error(run_id, error)
 
-    def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, Dict[str, Any], Dict[str, Any], UUID, Any) -> Any
+    def on_chain_start(
+        self: SentryLangchainCallback,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when chain starts running."""
         with capture_internal_exceptions():
             if not run_id:
@@ -326,8 +354,13 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
             if metadata:
                 set_data_normalized(watched_span.span, SPANDATA.AI_METADATA, metadata)
 
-    def on_chain_end(self, outputs, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, Dict[str, Any], UUID, Any) -> Any
+    def on_chain_end(
+        self: SentryLangchainCallback,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when chain ends running."""
         with capture_internal_exceptions():
             if not run_id or run_id not in self.span_map:
@@ -338,13 +371,23 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                 return
             self._exit_span(span_data, run_id)
 
-    def on_chain_error(self, error, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
+    def on_chain_error(
+        self: SentryLangchainCallback,
+        error: Union[Exception, KeyboardInterrupt],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when chain errors."""
         self._handle_error(run_id, error)
 
-    def on_agent_action(self, action, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, AgentAction, UUID, Any) -> Any
+    def on_agent_action(
+        self: SentryLangchainCallback,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         with capture_internal_exceptions():
             if not run_id:
                 return
@@ -360,8 +403,13 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                     watched_span.span, SPANDATA.AI_INPUT_MESSAGES, action.tool_input
                 )
 
-    def on_agent_finish(self, finish, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, AgentFinish, UUID, Any) -> Any
+    def on_agent_finish(
+        self: SentryLangchainCallback,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         with capture_internal_exceptions():
             if not run_id:
                 return
@@ -375,8 +423,14 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                 )
             self._exit_span(span_data, run_id)
 
-    def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, Dict[str, Any], str, UUID, Any) -> Any
+    def on_tool_start(
+        self: SentryLangchainCallback,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when tool starts running."""
         with capture_internal_exceptions():
             if not run_id:
@@ -399,8 +453,9 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                         watched_span.span, SPANDATA.AI_METADATA, kwargs.get("metadata")
                     )
 
-    def on_tool_end(self, output, *, run_id, **kwargs):
-        # type: (SentryLangchainCallback, str, UUID, Any) -> Any
+    def on_tool_end(
+        self: SentryLangchainCallback, output: str, *, run_id: UUID, **kwargs: Any
+    ) -> Any:
         """Run when tool ends running."""
         with capture_internal_exceptions():
             if not run_id or run_id not in self.span_map:
@@ -413,60 +468,95 @@ class SentryLangchainCallback(BaseCallbackHandler):  # type: ignore[misc]
                 set_data_normalized(span_data.span, SPANDATA.AI_RESPONSES, output)
             self._exit_span(span_data, run_id)
 
-    def on_tool_error(self, error, *args, run_id, **kwargs):
-        # type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
+    def on_tool_error(
+        self,
+        error: SentryLangchainCallback,
+        *args: Union[Exception, KeyboardInterrupt],
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
         """Run when tool errors."""
         self._handle_error(run_id, error)
 
 
-def _wrap_configure(f):
-    # type: (Callable[..., Any]) -> Callable[..., Any]
+def _wrap_configure(f: Callable[..., Any]) -> Callable[..., Any]:
 
     @wraps(f)
-    def new_configure(*args, **kwargs):
-        # type: (Any, Any) -> Any
+    def new_configure(
+        callback_manager_cls: type,
+        inheritable_callbacks: Callbacks = None,
+        local_callbacks: Callbacks = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
 
         integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
         if integration is None:
-            return f(*args, **kwargs)
+            return f(
+                callback_manager_cls,
+                inheritable_callbacks,
+                local_callbacks,
+                *args,
+                **kwargs,
+            )
 
-        with capture_internal_exceptions():
-            new_callbacks = []  # type: List[BaseCallbackHandler]
-            if "local_callbacks" in kwargs:
-                existing_callbacks = kwargs["local_callbacks"]
-                kwargs["local_callbacks"] = new_callbacks
-            elif len(args) > 2:
-                existing_callbacks = args[2]
-                args = (
-                    args[0],
-                    args[1],
-                    new_callbacks,
-                ) + args[3:]
-            else:
-                existing_callbacks = []
-
-            if existing_callbacks:
-                if isinstance(existing_callbacks, list):
-                    for cb in existing_callbacks:
-                        new_callbacks.append(cb)
-                elif isinstance(existing_callbacks, BaseCallbackHandler):
-                    new_callbacks.append(existing_callbacks)
-                else:
-                    logger.debug("Unknown callback type: %s", existing_callbacks)
-
-            already_added = False
-            for callback in new_callbacks:
-                if isinstance(callback, SentryLangchainCallback):
-                    already_added = True
-
-            if not already_added:
-                new_callbacks.append(
-                    SentryLangchainCallback(
-                        integration.max_spans,
-                        integration.include_prompts,
-                        integration.tiktoken_encoding_name,
-                    )
-                )
-        return f(*args, **kwargs)
+        local_callbacks = local_callbacks or []
+
+        # Handle each possible type of local_callbacks. For each type, we
+        # extract the list of callbacks to check for SentryLangchainCallback,
+        # and define a function that would add the SentryLangchainCallback
+        # to the existing callbacks list.
+        if isinstance(local_callbacks, BaseCallbackManager):
+            callbacks_list = local_callbacks.handlers
+        elif isinstance(local_callbacks, BaseCallbackHandler):
+            callbacks_list = [local_callbacks]
+        elif isinstance(local_callbacks, list):
+            callbacks_list = local_callbacks
+        else:
+            logger.debug("Unknown callback type: %s", local_callbacks)
+            # Just proceed with original function call
+            return f(
+                callback_manager_cls,
+                inheritable_callbacks,
+                local_callbacks,
+                *args,
+                **kwargs,
+            )
+
+        # Handle each possible type of inheritable_callbacks.
+        if isinstance(inheritable_callbacks, BaseCallbackManager):
+            inheritable_callbacks_list = inheritable_callbacks.handlers
+        elif isinstance(inheritable_callbacks, list):
+            inheritable_callbacks_list = inheritable_callbacks
+        else:
+            inheritable_callbacks_list = []
+
+        if not any(
+            isinstance(cb, SentryLangchainCallback)
+            for cb in itertools.chain(callbacks_list, inheritable_callbacks_list)
+        ):
+            sentry_handler = SentryLangchainCallback(
+                integration.max_spans,
+                integration.include_prompts,
+                integration.tiktoken_encoding_name,
+            )
+            if isinstance(local_callbacks, BaseCallbackManager):
+                local_callbacks = local_callbacks.copy()
+                local_callbacks.handlers = [
+                    *local_callbacks.handlers,
+                    sentry_handler,
+                ]
+            elif isinstance(local_callbacks, BaseCallbackHandler):
+                local_callbacks = [local_callbacks, sentry_handler]
+            else:  # local_callbacks is a list
+                local_callbacks = [*local_callbacks, sentry_handler]
+
+        return f(
+            callback_manager_cls,
+            inheritable_callbacks,
+            local_callbacks,
+            *args,
+            **kwargs,
+        )
 
     return new_configure