Migrate typing for integrations - part 3 (#4532)

Author: sl0thentr0py

sentry_sdk/integrations/langchain.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import itertools
 from collections import OrderedDict
 from functools import wraps
@@ -60,37 +61,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."""
 
-    def __init__(self, max_span_map_size, include_prompts, tiktoken_encoding_name=None):
-        # type: (int, bool, Optional[str]) -> None
-        self.span_map = OrderedDict()  # type: OrderedDict[UUID, WatchedSpan]
+    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
 
@@ -100,21 +105,18 @@ def __init__(self, max_span_map_size, include_prompts, tiktoken_encoding_name=No
 
             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
 
@@ -126,14 +128,17 @@ def _handle_error(self, run_id, error):
         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(
@@ -160,8 +165,9 @@ def _create_span(self, run_id, parent_id, **kwargs):
         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)
@@ -171,17 +177,16 @@ def _exit_span(self, span_data, run_id):
         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:
@@ -202,8 +207,14 @@ def on_llm_start(
                 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:
@@ -248,8 +259,9 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
                             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:
@@ -259,8 +271,13 @@ def on_llm_new_token(self, token, *, run_id, **kwargs):
                 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:
@@ -298,14 +315,25 @@ def on_llm_end(self, response, *, run_id, **kwargs):
 
             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:
@@ -325,8 +353,13 @@ def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
             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:
@@ -337,13 +370,23 @@ def on_chain_end(self, outputs, *, run_id, **kwargs):
                 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
@@ -359,8 +402,13 @@ def on_agent_action(self, action, *, run_id, **kwargs):
                     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
@@ -374,8 +422,14 @@ def on_agent_finish(self, finish, *, run_id, **kwargs):
                 )
             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:
@@ -398,8 +452,9 @@ def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
                         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:
@@ -412,24 +467,27 @@ def on_tool_end(self, output, *, run_id, **kwargs):
                 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(
-        callback_manager_cls,  # type: type
-        inheritable_callbacks=None,  # type: Callbacks
-        local_callbacks=None,  # type: Callbacks
-        *args,  # type: Any
-        **kwargs,  # type: Any
-    ):
-        # type: (...) -> Any
+        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: