add: patching for send_message_async

wandb · Oct 28, 2024 · a607896 · a607896
1 parent c04ca38
commit a607896
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/weave/integrations/google_ai_studio/google_ai_studio_sdk.py b/weave/integrations/google_ai_studio/google_ai_studio_sdk.py
@@ -110,5 +110,12 @@ async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
             "ChatSession.send_message",
             gemini_wrapper_sync(name="google.generativeai.ChatSession.send_message"),
         ),
+        SymbolPatcher(
+            lambda: importlib.import_module("google.generativeai.generative_models"),
+            "ChatSession.send_message_async",
+            gemini_wrapper_async(
+                name="google.generativeai.ChatSession.send_message_async"
+            ),
+        ),
     ]
 )