Use google-genai for chat

ag2ai · Jan 29, 2025 · 1b8af59 · 1b8af59
1 parent 2d38367
commit 1b8af59
Showing 1 changed file with 19 additions and 7 deletions.
diff --git a/autogen/oai/gemini.py b/autogen/oai/gemini.py
@@ -62,7 +62,7 @@
 from .client_utils import FormatterProtocol
 
 with optional_import_block():
-    import google.generativeai as genai
+    import google.genai as genai
     import vertexai
     from PIL import Image
     from google.auth.credentials import Credentials
@@ -71,6 +71,7 @@
         FunctionCall,
         FunctionDeclaration,
         FunctionResponse,
+        GenerateContentConfig,
         GenerateContentResponse,
         Part,
         Schema,
@@ -262,19 +263,30 @@ def create(self, params: dict) -> ChatCompletion:
             )
 
             chat = model.start_chat(history=gemini_messages[:-1], response_validation=response_validation)
+            response = chat.send_message(gemini_messages[-1].parts, stream=stream, safety_settings=safety_settings)
         else:
-            model = genai.GenerativeModel(
-                model_name,
-                generation_config=generation_config,
+            client = genai.Client(api_key=self.api_key)
+            generate_content_config = GenerateContentConfig(
                 safety_settings=safety_settings,
                 system_instruction=system_instruction,
                 tools=tools,
+                **generation_config,
             )
+            chat = client.chats.create(model=model_name, config=generate_content_config, history=gemini_messages[:-1])
+            response = chat.send_message(message=gemini_messages[-1].parts)
+
+            # model = genai.GenerativeModel(
+            #     model_name,
+            #     generation_config=generation_config,
+            #     safety_settings=safety_settings,
+            #     system_instruction=system_instruction,
+            #     tools=tools,
+            # )
 
-            genai.configure(api_key=self.api_key)
-            chat = model.start_chat(history=gemini_messages[:-1])
+            # genai.configure(api_key=self.api_key)
+            # chat = model.start_chat(history=gemini_messages[:-1])
 
-        response = chat.send_message(gemini_messages[-1].parts, stream=stream, safety_settings=safety_settings)
+            # response = chat.send_message(gemini_messages[-1].parts, stream=stream, safety_settings=safety_settings)
 
         # Extract text and tools from response
         ans = ""