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

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

sentry_sdk/integrations/launchdarkly.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from sentry_sdk.feature_flags import add_feature_flag
@@ -20,8 +21,7 @@ except ImportError:
 class LaunchDarklyIntegration(Integration):
     identifier = "launchdarkly"
 
-    def __init__(self, ld_client=None):
-        # type: (LDClient | None) -> None
+    def __init__(self, ld_client: LDClient | None = None) -> None:
         """
         :param client: An initialized LDClient instance. If a client is not provided, this
             integration will attempt to use the shared global instance.
@@ -38,25 +38,28 @@ class LaunchDarklyIntegration(Integration):
         client.add_hook(LaunchDarklyHook())
 
     @staticmethod
-    def setup_once():
-        # type: () -> None
+    def setup_once() -> None:
         pass
 
 
 class LaunchDarklyHook(Hook):
 
     @property
-    def metadata(self):
-        # type: () -> Metadata
+    def metadata(self) -> Metadata:
         return Metadata(name="sentry-flag-auditor")
 
-    def after_evaluation(self, series_context, data, detail):
-        # type: (EvaluationSeriesContext, dict[Any, Any], EvaluationDetail) -> dict[Any, Any]
+    def after_evaluation(
+        self,
+        series_context: EvaluationSeriesContext,
+        data: dict[Any, Any],
+        detail: EvaluationDetail,
+    ) -> dict[Any, Any]:
         if isinstance(detail.value, bool):
             add_feature_flag(series_context.key, detail.value)
 
         return data
 
-    def before_evaluation(self, series_context, data):
-        # type: (EvaluationSeriesContext, dict[Any, Any]) -> dict[Any, Any]
+    def before_evaluation(
+        self, series_context: EvaluationSeriesContext, data: dict[Any, Any]
+    ) -> dict[Any, Any]:
         return data  # No-op.