Spaces:
Running
Running
| import json | |
| import logging | |
| import pandas as pd | |
| from typing import Dict, Any, Tuple | |
| from core.llm import chat | |
| logger = logging.getLogger(__name__) | |
| def parse_what_if_query(query: str, columns: list) -> Dict[str, Any]: | |
| """ | |
| Uses an LLM to extract the causal simulation parameters from a What-If query. | |
| """ | |
| prompt = f"""You are a Causal AI Simulator. | |
| The user is asking a "What-If" question to simulate a change in their dataset. | |
| AVAILABLE COLUMNS IN DATASET: | |
| {columns} | |
| USER QUERY: "{query}" | |
| Extract the simulation parameters. | |
| 1. 'manipulated_variable': The exact column name from the list above that the user wants to change. | |
| 2. 'operation': 'increase', 'decrease', 'set', or 'multiply'. | |
| 3. 'value': The numerical value of the change (e.g. 20 for 20%). | |
| 4. 'is_percentage': true if they said "20%", false if they said "by $20". | |
| Output ONLY a valid JSON object: | |
| {{ | |
| "manipulated_variable": "Column_Name", | |
| "operation": "increase", | |
| "value": 20, | |
| "is_percentage": true | |
| }} | |
| If no simulation parameters can be parsed, return an empty JSON {{}}. | |
| """ | |
| try: | |
| response = chat(prompt, temperature=0.1, max_tokens=150) | |
| # Parse JSON | |
| response = response.strip() | |
| if '```json' in response: | |
| response = response.split('```json')[1].split('```')[0] | |
| elif '```' in response: | |
| response = response.split('```')[1].split('```')[0] | |
| start = response.find('{') | |
| end = response.rfind('}') + 1 | |
| if start >= 0 and end > start: | |
| response = response[start:end] | |
| decision = json.loads(response) | |
| return decision | |
| except Exception as e: | |
| logger.error(f"Failed to parse what-if query: {e}") | |
| return {} | |
| def simulate_what_if(query: str, df: pd.DataFrame) -> Tuple[bool, pd.DataFrame, Dict[str, Any]]: | |
| """ | |
| Simulates a what-if scenario on the DataFrame based on the user's query. | |
| Returns (Success: bool, Simulated_DF: DataFrame, Simulation_Details: Dict) | |
| """ | |
| if df is None or df.empty: | |
| return False, df, {} | |
| try: | |
| columns = list(df.columns) | |
| sim_params = parse_what_if_query(query, columns) | |
| if not sim_params or "manipulated_variable" not in sim_params: | |
| return False, df, {} | |
| col = sim_params["manipulated_variable"] | |
| op = sim_params.get("operation", "increase") | |
| val = float(sim_params.get("value", 0)) | |
| is_pct = sim_params.get("is_percentage", False) | |
| if col not in df.columns: | |
| return False, df, {"error": f"Column {col} not found"} | |
| # Create a deep copy of the dataframe for simulation | |
| sim_df = df.copy() | |
| # Ensure column is numeric before manipulating | |
| if pd.api.types.is_numeric_dtype(sim_df[col]): | |
| if op == "increase": | |
| if is_pct: | |
| sim_df[col] = sim_df[col] * (1 + (val / 100)) | |
| else: | |
| sim_df[col] = sim_df[col] + val | |
| elif op == "decrease": | |
| if is_pct: | |
| sim_df[col] = sim_df[col] * (1 - (val / 100)) | |
| else: | |
| sim_df[col] = sim_df[col] - val | |
| elif op == "multiply": | |
| sim_df[col] = sim_df[col] * val | |
| elif op == "set": | |
| sim_df[col] = val | |
| return True, sim_df, sim_params | |
| else: | |
| return False, df, {"error": f"Column {col} is not numeric."} | |
| except Exception as e: | |
| logger.error(f"Error in simulate_what_if: {e}") | |
| return False, df, {"error": str(e)} | |