Update insight_and_tasks/utils/pandasai_setup.py

insight_and_tasks/utils/pandasai_setup.py

@@ -25,28 +25,22 @@ def configure_pandasai(api_key: str, model_name: str = None):
         # raise ValueError("API key must be provided for PandasAI configuration")
         return
 
+    os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "False"
+    os.environ["GOOGLE_API_KEY"] = api_key
+
     selected_model = model_name if model_name else DEFAULT_PANDASAI_MODEL
     
     try:
         llm = LiteLLM(
             model=selected_model, # Use the selected model
             api_key=api_key
-            # You might need to add other parameters for LiteLLM depending on the provider
-            # e.g., if not using a Google model directly via gemini provider in LiteLLM
         )
 
         # PandasAI configuration
         pai.config.set({
             "llm": llm,
-            "verbose": os.environ.get("PANDASAI_VERBOSE", "False").lower() == "true",
-            "enable_cache": True,
-            "enforce_privacy": False, # Be cautious with this in production
-            "save_charts": False, # Set to True if you want to save charts locally
-            # "save_charts_path": "charts_output", # Define path if saving charts
-            "custom_whitelisted_dependencies": [], # Add any custom dependencies if needed
-            "max_retries": 3, # Default retries for PandasAI operations
-            "temperature": 0.3, # Lower temperature for more deterministic/factual outputs
-            # "open_charts": False # Whether to automatically open charts
+            "temperature": 0.3,  # Lower temperature for more consistent results
+            "max_retries": 3
         })
         logger.info(f"PandasAI configured successfully with model: {selected_model}")
         logger.info(f"PandasAI LLM object: {llm}")