feat(litellm): Add async callbacks

scripts/populate_tox/config.py

@@ -238,7 +238,7 @@
     "litellm": {
         "package": "litellm",
         "deps": {
-            "*": ["anthropic", "google-genai"],
+            "*": ["anthropic", "google-genai", "pytest-asyncio"],
         },
     },
     "litestar": {

sentry_sdk/integrations/litellm.py

@@ -82,7 +82,7 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None:
         provider = "unknown"
 
     call_type = kwargs.get("call_type", None)
-    if call_type == "embedding":
+    if call_type == "embedding" or call_type == "aembedding":
         operation = "embeddings"
     else:
         operation = "chat"
@@ -170,6 +170,10 @@ def _input_callback(kwargs: "Dict[str, Any]") -> None:
             set_data_normalized(span, f"gen_ai.litellm.{key}", value)
 
 
+async def _async_input_callback(kwargs: "Dict[str, Any]") -> None:
+    return _input_callback(kwargs)
+
+
 def _success_callback(
     kwargs: "Dict[str, Any]",
     completion_response: "Any",
@@ -234,12 +238,30 @@ def _success_callback(
         is_streaming = kwargs.get("stream")
         # Callback is fired multiple times when streaming a response.
         # Streaming flag checked at https://github.com/BerriAI/litellm/blob/33c3f13443eaf990ac8c6e3da78bddbc2b7d0e7a/litellm/litellm_core_utils/litellm_logging.py#L1603
-        if is_streaming is not True or "complete_streaming_response" in kwargs:
+        if (
+            is_streaming is not True
+            or "complete_streaming_response" in kwargs
+            or "async_complete_streaming_response" in kwargs
+        ):
             span = metadata.pop("_sentry_span", None)
             if span is not None:
                 span.__exit__(None, None, None)
 
 
+async def _async_success_callback(
+    kwargs: "Dict[str, Any]",
+    completion_response: "Any",
+    start_time: "datetime",
+    end_time: "datetime",
+) -> None:
+    return _success_callback(
+        kwargs,
+        completion_response,
+        start_time,
+        end_time,
+    )
+
+
 def _failure_callback(
     kwargs: "Dict[str, Any]",
     exception: Exception,
@@ -321,10 +343,14 @@ def setup_once() -> None:
         litellm.input_callback = input_callback or []
         if _input_callback not in litellm.input_callback:
             litellm.input_callback.append(_input_callback)
+        if _async_input_callback not in litellm.input_callback:
+            litellm.input_callback.append(_async_input_callback)
 
         litellm.success_callback = success_callback or []
         if _success_callback not in litellm.success_callback:
             litellm.success_callback.append(_success_callback)
+        if _async_success_callback not in litellm.success_callback:
+            litellm.success_callback.append(_async_success_callback)
 
         litellm.failure_callback = failure_callback or []
         if _failure_callback not in litellm.failure_callback: