Merge remote-tracking branch 'origin/master' into potel-base

Author: sl0thentr0py

CHANGELOG.md

@@ -40,6 +40,16 @@ for your feedback. How was the migration? Is everything working as expected? Is
 [on GitHub](https://github.com/getsentry/sentry-python/discussions/3936) or
 [on Discord](https://discord.com/invite/Ww9hbqr).
 
+## 2.33.0
+
+### Various fixes & improvements
+
+- feat(langchain): Support `BaseCallbackManager` (#4486) by @szokeasaurusrex
+- Use `span.data` instead of `measurements` for token usage (#4567) by @antonpirker
+- Fix custom model name (#4569) by @antonpirker
+- fix: shut down "session flusher" more promptly (#4561) by @bukzor
+- chore: Remove Lambda urllib3 pin on Python 3.10+ (#4549) by @sentrivana
+
 ## 2.32.0
 
 ### Various fixes & improvements

sentry_sdk/ai/monitoring.py

@@ -96,26 +96,37 @@ async def async_wrapped(*args: Any, **kwargs: Any) -> Any:
 
 def record_token_usage(
     span: Span,
-    prompt_tokens: Optional[int] = None,
-    completion_tokens: Optional[int] = None,
+    input_tokens: Optional[int] = None,
+    input_tokens_cached: Optional[int] = None,
+    output_tokens: Optional[int] = None,
+    output_tokens_reasoning: Optional[int] = None,
     total_tokens: Optional[int] = None,
 ) -> None:
+    # TODO: move pipeline name elsewhere
     ai_pipeline_name = get_ai_pipeline_name()
     if ai_pipeline_name:
         span.set_attribute(SPANDATA.AI_PIPELINE_NAME, ai_pipeline_name)
 
-    if prompt_tokens is not None:
-        span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens)
+    if input_tokens is not None:
+        span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
 
-    if completion_tokens is not None:
-        span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens)
+    if input_tokens_cached is not None:
+        span.set_attribute(
+            SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED,
+            input_tokens_cached,
+        )
 
-    if (
-        total_tokens is None
-        and prompt_tokens is not None
-        and completion_tokens is not None
-    ):
-        total_tokens = prompt_tokens + completion_tokens
+    if output_tokens is not None:
+        span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
+
+    if output_tokens_reasoning is not None:
+        span.set_data(
+            SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING,
+            output_tokens_reasoning,
+        )
+
+    if total_tokens is None and input_tokens is not None and output_tokens is not None:
+        total_tokens = input_tokens + output_tokens
 
     if total_tokens is not None:
         span.set_attribute(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)

sentry_sdk/integrations/anthropic.py

@@ -62,7 +62,13 @@ def _calculate_token_usage(result: Messages, span: Span) -> None:
             output_tokens = usage.output_tokens
 
     total_tokens = input_tokens + output_tokens
-    record_token_usage(span, input_tokens, output_tokens, total_tokens)
+
+    record_token_usage(
+        span,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=total_tokens,
+    )
 
 
 def _get_responses(content: list[Any]) -> list[dict[str, Any]]:
@@ -129,7 +135,12 @@ def _add_ai_data_to_span(
                 [{"type": "text", "text": complete_message}],
             )
         total_tokens = input_tokens + output_tokens
-        record_token_usage(span, input_tokens, output_tokens, total_tokens)
+        record_token_usage(
+            span,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+        )
         span.set_attribute(SPANDATA.AI_STREAMING, True)
 
 

sentry_sdk/integrations/cohere.py

@@ -114,14 +114,14 @@ def collect_chat_response_fields(
             if hasattr(res.meta, "billed_units"):
                 record_token_usage(
                     span,
-                    prompt_tokens=res.meta.billed_units.input_tokens,
-                    completion_tokens=res.meta.billed_units.output_tokens,
+                    input_tokens=res.meta.billed_units.input_tokens,
+                    output_tokens=res.meta.billed_units.output_tokens,
                 )
             elif hasattr(res.meta, "tokens"):
                 record_token_usage(
                     span,
-                    prompt_tokens=res.meta.tokens.input_tokens,
-                    completion_tokens=res.meta.tokens.output_tokens,
+                    input_tokens=res.meta.tokens.input_tokens,
+                    output_tokens=res.meta.tokens.output_tokens,
                 )
 
             if hasattr(res.meta, "warnings"):
@@ -258,7 +258,7 @@ def new_embed(*args: Any, **kwargs: Any) -> Any:
             ):
                 record_token_usage(
                     span,
-                    prompt_tokens=res.meta.billed_units.input_tokens,
+                    input_tokens=res.meta.billed_units.input_tokens,
                     total_tokens=res.meta.billed_units.input_tokens,
                 )
             return res

sentry_sdk/integrations/huggingface_hub.py

@@ -108,7 +108,10 @@ def new_text_generation(*args: Any, **kwargs: Any) -> Any:
                         [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
 
@@ -141,7 +144,10 @@ def new_details_iterator() -> Iterable[ChatCompletionStreamOutput]:
                                 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()

sentry_sdk/integrations/langchain.py

@@ -24,6 +24,7 @@
     from langchain_core.callbacks import (
         manager,
         BaseCallbackHandler,
+        BaseCallbackManager,
         Callbacks,
     )
     from langchain_core.agents import AgentAction, AgentFinish
@@ -302,15 +303,15 @@ def on_llm_end(
                 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)
@@ -499,12 +500,20 @@ def new_configure(
                 **kwargs,
             )
 
-        callbacks_list = local_callbacks or []
-
-        if isinstance(callbacks_list, BaseCallbackHandler):
-            callbacks_list = [callbacks_list]
-        elif not isinstance(callbacks_list, list):
-            logger.debug("Unknown callback type: %s", callbacks_list)
+        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,
@@ -514,28 +523,38 @@ def new_configure(
                 **kwargs,
             )
 
-        inheritable_callbacks_list = (
-            inheritable_callbacks if isinstance(inheritable_callbacks, list) else []
-        )
+        # 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)
         ):
-            # Avoid mutating the existing callbacks list
-            callbacks_list = [
-                *callbacks_list,
-                SentryLangchainCallback(
-                    integration.max_spans,
-                    integration.include_prompts,
-                    integration.tiktoken_encoding_name,
-                ),
-            ]
+            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,
-            callbacks_list,
+            local_callbacks,
             *args,
             **kwargs,
         )