Update

agents.py

@@ -5,18 +5,14 @@ import json
 
 logger = logging.getLogger(__name__)
 
-# Try both import methods with proper error handling and model access
+# Try both import methods with proper error handling
 try:
     import google.generativeai as genai
     GENAI_PACKAGE = "generativeai"
-    MODEL_CLASS = genai.GenerativeModel
 except ImportError:
     try:
         import google.genai as genai
         GENAI_PACKAGE = "genai"
-        def create_model(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")
         raise
@@ -31,14 +27,13 @@ class BaseAgent:
             if not api_key:
                 raise ValueError("Gemini API key is required when use_gemini=True")
             
-            try:
-                if GENAI_PACKAGE == "generativeai":
+            # Set API key for either package version
+            os.environ["GOOGLE_API_KEY"] = api_key
+            if GENAI_PACKAGE == "generativeai":
+                try:
                     genai.configure(api_key=api_key)
-                else:
-                    os.environ["GOOGLE_API_KEY"] = api_key
-            except Exception as e:
-                logger.warning(f"Using fallback configuration method: {str(e)}")
-                os.environ["GOOGLE_API_KEY"] = api_key
+                except Exception as e:
+                    logger.warning(f"Fallback to environment variable: {str(e)}")
                 
             self.gemini_model = gemini_model or "gemini-1.5-pro"
         else:
@@ -50,30 +45,23 @@ class BaseAgent:
 
     def _generate_with_gemini(self, prompt: str, system_prompt: str) -> str:
         try:
+            combined_prompt = f"System: {system_prompt}\n\nUser: {prompt}"
+            
             if GENAI_PACKAGE == "generativeai":
-                model = MODEL_CLASS(model_name=self.gemini_model)
+                model = genai.GenerativeModel(model_name=self.gemini_model)
+                response = model.generate_content(
+                    combined_prompt,
+                    generation_config={"temperature": 0.1}
+                )
             else:
-                model = MODEL_CLASS(self.gemini_model)
-                
-            combined_prompt = f"System: {system_prompt}\n\nUser: {prompt}"
+                # Use basic text generation for google.genai
+                completion = genai.generate_text(
+                    model=self.gemini_model,
+                    prompt=combined_prompt
+                )
+                return str(completion)
             
-            try:
-                if GENAI_PACKAGE == "generativeai":
-                    response = model.generate_content(
-                        combined_prompt,
-                        generation_config={"temperature": 0.1}
-                    )
-                else:
-                    response = model.generate_text(combined_prompt)
-                
-                return response.text if hasattr(response, 'text') else str(response)
-            except AttributeError:
-                # Fallback for different response formats
-                if GENAI_PACKAGE == "generativeai":
-                    response = model.generate_content(combined_prompt)
-                else:
-                    response = model.generate_text(combined_prompt)
-                return response.text if hasattr(response, 'text') else str(response)
+            return response.text if hasattr(response, 'text') else str(response)
                 
         except Exception as e:
             logger.error(f"Gemini generation failed: {str(e)}")