|
|
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) |
|
|
|
|
|
|
|
|
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": [] |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
