From a607896bfefcc715ecaa955d06a62b4722f6fa35 Mon Sep 17 00:00:00 2001 From: Soumik Rakshit <19soumik.rakshit96@gmail.com> Date: Tue, 29 Oct 2024 03:23:53 +0530 Subject: [PATCH] add: patching for send_message_async --- .../integrations/google_ai_studio/google_ai_studio_sdk.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/weave/integrations/google_ai_studio/google_ai_studio_sdk.py b/weave/integrations/google_ai_studio/google_ai_studio_sdk.py index 57045942282..91cb31b8130 100644 --- 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" + ), + ), ] )