From cf2f08a08e2f273b8e241f2095fbb471659e3bf1 Mon Sep 17 00:00:00 2001 From: Soumik Rakshit <19soumik.rakshit96@gmail.com> Date: Tue, 6 Aug 2024 13:19:38 +0530 Subject: [PATCH] add: support for generate_content_async and start_chat --- weave/integrations/gemini/gemini_sdk.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/weave/integrations/gemini/gemini_sdk.py b/weave/integrations/gemini/gemini_sdk.py index 8a0a66db793..a3523f50d18 100644 --- a/weave/integrations/gemini/gemini_sdk.py +++ b/weave/integrations/gemini/gemini_sdk.py @@ -50,21 +50,24 @@ def wrapper(fn: Callable) -> Callable: gemini_wrapper(name="google.generativeai.GenerativeModel.generate_content"), ), SymbolPatcher( - lambda: importlib.import_module( - "google.generativeai.types.generation_types" - ), - "GenerateContentResponse.from_response", + lambda: importlib.import_module("google.generativeai"), + "GenerativeModel.generate_content_async", gemini_wrapper( - name="google.generativeai.types.generation_types.GenerateContentResponse.from_response" + name="google.generativeai.GenerativeModel.generate_content_async" ), ), + SymbolPatcher( + lambda: importlib.import_module("google.generativeai"), + "GenerativeModel.generate_content", + gemini_wrapper(name="google.generativeai.GenerativeModel.start_chat"), + ), SymbolPatcher( lambda: importlib.import_module( "google.generativeai.types.generation_types" ), - "GenerateContentResponse.from_iterator", + "GenerateContentResponse.from_response", gemini_wrapper( - name="google.generativeai.types.generation_types.GenerateContentResponse.from_iterator" + name="google.generativeai.types.generation_types.GenerateContentResponse.from_response" ), ), SymbolPatcher(