Spaces:
Running
Running
| # 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 | |
| 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=DEFAULT_PANDASAI_MODEL , # Use the selected model | |
| api_key=api_key | |
| ) | |
| # PandasAI configuration | |
| pai.config.set({ | |
| "llm": llm, | |
| "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}") | |
| 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) | |