Spaces:
Running
Running
File size: 1,904 Bytes
40067b3 0b3173b cd6ad96 40067b3 12bbfb7 40067b3 0b3173b 40067b3 a985715 40067b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
# 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)
|