Update

agents.py

@@ -15,7 +15,7 @@ except ImportError:
         import google.genai as genai
         GENAI_PACKAGE = "genai"
         def create_model(model_name):
-            return genai.models[model_name]
+            return genai.get_model(model_name)  # Using get_model instead of accessing models directly
         MODEL_CLASS = create_model
     except ImportError:
         logger.error("Failed to import Google AI package")
@@ -50,7 +50,11 @@ class BaseAgent:
 
     def _generate_with_gemini(self, prompt: str, system_prompt: str) -> str:
         try:
-            model = MODEL_CLASS(self.gemini_model)
+            if GENAI_PACKAGE == "generativeai":
+                model = MODEL_CLASS(model_name=self.gemini_model)
+            else:
+                model = MODEL_CLASS(self.gemini_model)
+                
             combined_prompt = f"System: {system_prompt}\n\nUser: {prompt}"
             
             try: