Create pandasai_setup.py

insight_and_tasks/utils/pandasai_setup.py

@@ -0,0 +1,59 @@
+# utils/pandasai_setup.py
+import os
+import logging
+import pandasai as pai
+from pandasai_litellm import LiteLLM # Ensure this import matches your installed library
+
+# Configure logger for this module
+logger = logging.getLogger(__name__)
+
+# It's good practice to define constants at the top or in a config file
+DEFAULT_PANDASAI_MODEL = "gemini/gemini-2.5-flash-preview-05-20" # Using a common default
+
+def configure_pandasai(api_key: str, model_name: str = None):
+    """
+    Configures PandasAI with LiteLLM using the provided API key and model.
+
+    Args:
+        api_key: The Google API key.
+        model_name: The specific model to use (e.g., "gemini/gemini-1.5-flash-latest"). 
+                    If None, uses DEFAULT_PANDASAI_MODEL.
+    """
+    if not api_key:
+        logger.error("PandasAI Configuration Error: API key is missing.")
+        # Depending on strictness, you might raise an error or just log
+        # raise ValueError("API key must be provided for PandasAI configuration")
+        return
+
+    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
+        })
+        logger.info(f"PandasAI configured successfully with model: {selected_model}")
+        logger.info(f"PandasAI LLM object: {pai.config.llm}")
+
+
+    except ImportError:
+        logger.error("PandasAI or pandasai_litellm is not installed. Please install the required packages.")
+    except Exception as e:
+        logger.error(f"Error configuring PandasAI: {e}", exc_info=True)
+