import json import numpy as np import pandas as pd import re import os import uuid import logging from io import StringIO import sys import traceback from typing import Optional, Dict, Any, List from pydantic import BaseModel, Field from google.generativeai import GenerativeModel, configure from dotenv import load_dotenv import seaborn as sns import datetime as dt pd.set_option('display.max_columns', None) pd.set_option('display.max_rows', None) pd.set_option('display.max_colwidth', None) load_dotenv() API_KEYS = os.getenv("GEMINI_API_KEYS", "").split(",")[::-1] MODEL_NAME = 'gemini-2.0-flash' os.environ['MPLBACKEND'] = 'agg' import matplotlib.pyplot as plt plt.show = lambda: None logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) class GeminiKeyManager: """Manage multiple Gemini API keys with failover""" def __init__(self, api_keys: List[str]): self.original_keys = api_keys.copy() self.available_keys = api_keys.copy() self.active_key = None self.failed_keys = {} def configure(self) -> bool: while self.available_keys: key = self.available_keys.pop(0) try: configure(api_key=key) self.active_key = key logger.info(f"Configured with key: {self._mask_key(key)}") return True except Exception as e: self.failed_keys[key] = str(e) logger.error(f"Key failed: {self._mask_key(key)}. Error: {str(e)}") logger.critical("All API keys failed") return False def _mask_key(self, key: str) -> str: return f"{key[:8]}...{key[-4:]}" if key else "" class PythonREPL: """Secure Python REPL with file generation tracking""" def __init__(self, df: pd.DataFrame): self.df = df self.output_dir = os.path.abspath(f'generated_outputs/{uuid.uuid4()}') os.makedirs(self.output_dir, exist_ok=True) self.local_env = { "pd": pd, "df": self.df.copy(), "plt": plt, "os": os, "uuid": uuid, "sns": sns, "json": json, "dt": dt, "output_dir": self.output_dir } def execute(self, code: str) -> Dict[str, Any]: print('Executing code...', code) old_stdout = sys.stdout sys.stdout = mystdout = StringIO() file_tracker = { 'csv_files': set(), 'image_files': set() } try: code = f""" import matplotlib.pyplot as plt plt.switch_backend('agg') {code} plt.close('all') """ exec(code, self.local_env) self.df = self.local_env.get('df', self.df) # Track generated files for fname in os.listdir(self.output_dir): if fname.endswith('.csv'): file_tracker['csv_files'].add(fname) elif fname.lower().endswith(('.png', '.jpg', '.jpeg')): file_tracker['image_files'].add(fname) error = False except Exception as e: error_msg = traceback.format_exc() error = True finally: sys.stdout = old_stdout return { "output": mystdout.getvalue(), "error": error, "error_message": error_msg if error else None, "df": self.local_env.get('df', self.df), "output_dir": self.output_dir, "files": { "csv": [os.path.join(self.output_dir, f) for f in file_tracker['csv_files']], "images": [os.path.join(self.output_dir, f) for f in file_tracker['image_files']] } } class RethinkAgent(BaseModel): df: pd.DataFrame max_retries: int = Field(default=5, ge=1) gemini_model: Optional[GenerativeModel] = None current_retry: int = Field(default=0, ge=0) repl: Optional[PythonREPL] = None key_manager: Optional[GeminiKeyManager] = None class Config: arbitrary_types_allowed = True def _extract_code(self, response: str) -> str: code_match = re.search(r'```python(.*?)```', response, re.DOTALL) return code_match.group(1).strip() if code_match else response.strip() def _generate_initial_prompt(self, query: str) -> str: return f"""Generate DIRECT EXECUTION CODE (no functions, no explanations) following STRICT RULES: MANDATORY REQUIREMENTS: 1. Operate directly on existing 'df' variable 2. Save ALL final DataFrames to CSV using: df.to_csv(f'{{output_dir}}/descriptive_name.csv') 3. For visualizations: plt.savefig(f'{{output_dir}}/chart_name.png') 4. Use EXACTLY this structure: # Data processing df_processed = df[...] # filtering/grouping # Save results df_processed.to_csv(f'{{output_dir}}/result.csv') # Visualizations (if needed) plt.figure() ... plotting code ... plt.savefig(f'{{output_dir}}/chart.png') plt.close() FORBIDDEN: - Function definitions - Dummy data creation - Any code blocks besides pandas operations and matplotlib - Print statements showing dataframes DATAFRAME COLUMNS: {', '.join(self.df.columns)} USER QUERY: {query} EXAMPLE RESPONSE FOR "Sales by region": # Data processing sales_by_region = df.groupby('region')['sales'].sum().reset_index() # Save results sales_by_region.to_csv(f'{{output_dir}}/sales_by_region.csv') """ def _generate_retry_prompt(self, query: str, error: str, code: str) -> str: return f"""FIX THIS CODE (failed with: {error}) by STRICTLY FOLLOWING: 1. REMOVE ALL FUNCTION DEFINITIONS 2. ENSURE DIRECT DF OPERATIONS 3. USE EXPLICIT output_dir PATHS 4. ADD NECESSARY IMPORTS IF MISSING 5. VALIDATE COLUMN NAMES EXIST BAD CODE: {code} CORRECTED CODE:""" def initialize_model(self, api_keys: List[str]) -> bool: self.key_manager = GeminiKeyManager(api_keys) if not self.key_manager.configure(): raise RuntimeError("API key initialization failed") try: self.gemini_model = GenerativeModel(MODEL_NAME) return True except Exception as e: logger.error(f"Model init failed: {str(e)}") return False def generate_code(self, query: str, error: Optional[str] = None, previous_code: Optional[str] = None) -> str: prompt = self._generate_retry_prompt(query, error, previous_code) if error else self._generate_initial_prompt(query) try: response = self.gemini_model.generate_content(prompt) return self._extract_code(response.text) except Exception as e: if self.key_manager.available_keys and self.key_manager.configure(): return self.generate_code(query, error, previous_code) raise def execute_query(self, query: str) -> Dict[str, Any]: self.repl = PythonREPL(self.df) result = None while self.current_retry < self.max_retries: try: code = self.generate_code(query, result["error_message"] if result else None, result["code"] if result else None) execution_result = self.repl.execute(code) if execution_result["error"]: self.current_retry += 1 result = { "error_message": execution_result["error_message"], "code": code } else: return { "text": execution_result["output"], "csv_files": execution_result["files"]["csv"], "image_files": execution_result["files"]["images"] } except Exception as e: return { "error": f"Critical failure: {str(e)}", "csv_files": [], "image_files": [] } return { "error": f"Failed after {self.max_retries} retries", "csv_files": [], "image_files": [] } def gemini_llm_chat(csv_url: str, query: str) -> Dict[str, Any]: try: df = pd.read_csv(csv_url) agent = RethinkAgent(df=df) if not agent.initialize_model(API_KEYS): return {"error": "API configuration failed"} result = agent.execute_query(query) if "error" in result: return result return { "message": result["text"], "csv_files": result["csv_files"], "image_files": result["image_files"] } except Exception as e: logger.error(f"Processing failed: {str(e)}") return { "error": f"Processing error: {str(e)}", "csv_files": [], "image_files": [] } def generate_csv_report(csv_url: str, query: str): try: result = gemini_llm_chat(csv_url, query) json_result = json.dumps(result, indent=2) logger.info(f"Report generated successfully: {json_result}") return json_result except Exception as e: logger.error(f"Report generation failed: {str(e)}") return { "error": f"Report generation error: {str(e)}", "csv_files": [], "image_files": [] } # if __name__ == "__main__": # result = gemini_llm_chat("./documents/enterprise_sales_data.csv", # "Generate a detailed sales report of the last 6 months from all the aspects and include a bar chart showing the sales by region.") # print(json.dumps(result, indent=2))