Update app.py

app.py

@@ -1,12 +1,13 @@
 # Odyssey - The AI Data Science Workspace
-# A demonstration of a state-of-the-art, AI-native analytic environment.
+# A state-of-the-art, AI-native analytic environment.
+# This script is a complete, self-contained Gradio application.
 
 import gradio as gr
 import pandas as pd
 import numpy as np
 import plotly.express as px
 import plotly.graph_objects as go
-import io, os, json, base64, logging, warnings, pickle, uuid
+import io, os, json, pickle, logging, warnings, uuid
 from contextlib import redirect_stdout
 from datetime import datetime
 
@@ -18,13 +19,20 @@ from sklearn.metrics import roc_curve, auc, confusion_matrix, r2_score, mean_squ
 from sklearn.preprocessing import LabelEncoder
 from sklearn.impute import KNNImputer
 
+# Optional: For AI features
+try:
+    import google.generativeai as genai
+except ImportError:
+    print("Warning: 'google-generativeai' not found. AI features will be disabled.")
+    genai = None
+
 # --- Configuration ---
 warnings.filterwarnings('ignore')
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # --- UI Theme & Icons ---
 THEME = gr.themes.Monochrome(primary_hue="indigo", secondary_hue="blue", neutral_hue="slate").set(
-    body_background_fill="radial-gradient(circle, rgba(20,20,80,1) 0%, rgba(0,0,10,1) 100%);",
+    body_background_fill="radial-gradient(circle, rgba(10,20,50,1) 0%, rgba(0,0,10,1) 100%);",
     block_label_background_fill="rgba(255,255,255,0.05)",
     block_background_fill="rgba(255,255,255,0.05)",
     button_primary_background_fill="linear-gradient(90deg, #6A11CB 0%, #2575FC 100%)",
@@ -33,57 +41,49 @@ THEME = gr.themes.Monochrome(primary_hue="indigo", secondary_hue="blue", neutral
 )
 ICONS = {"overview": "🔭", "medic": "🧪", "launchpad": "🚀", "copilot": "💡", "export": "📄"}
 
+# --- Helper Functions ---
+def safe_exec(code_string: str, local_vars: dict) -> tuple:
+    """Safely execute a string of Python code and capture its output."""
+    output_buffer = io.StringIO()
+    try:
+        with redirect_stdout(output_buffer):
+            exec(code_string, globals(), local_vars)
+        stdout = output_buffer.getvalue()
+        fig = local_vars.get('fig')
+        df_out = local_vars.get('df_result')
+        return stdout, fig, df_out, None
+    except Exception as e:
+        return None, None, None, f"Execution Error: {str(e)}"
+
 # --- Core State & Project Management ---
 def init_state():
-    """Initializes a blank global state."""
+    """Initializes a blank global state dictionary."""
     return {
-        "project_name": None,
-        "df_original": None,
-        "df_modified": None,
-        "metadata": None,
-        "insights": None,
-        "chat_history": [],
-        "dynamic_dashboards": {}
+        "project_name": None, "df_original": None, "df_modified": None,
+        "metadata": None, "insights": None, "chat_history": []
     }
 
 def save_project(state):
-    """Saves the entire application state to a .osyssey file."""
+    """Saves the entire application state to a .odyssey file."""
     if not state or not state.get("project_name"):
         return gr.update(value="Project needs a name to save.", interactive=True)
     
     filename = f"{state['project_name']}.odyssey"
-    # Convert dataframes to pickle strings for serialization
-    state_to_save = state.copy()
-    if state_to_save['df_original'] is not None:
-        state_to_save['df_original'] = state_to_save['df_original'].to_pickle()
-    if state_to_save['df_modified'] is not None:
-        state_to_save['df_modified'] = state_to_save['df_modified'].to_pickle()
-
     with open(filename, "wb") as f:
-        pickle.dump(state_to_save, f)
-    
-    return gr.update(value=f"Project saved to {filename}", interactive=True)
+        pickle.dump(state, f)
+    return gr.update(value=f"✅ Project saved to {filename}", interactive=True)
 
 def load_project(file_obj):
     """Loads a .odyssey file into the application state."""
     if not file_obj: return init_state()
     with open(file_obj.name, "rb") as f:
-        loaded_state = pickle.load(f)
-    
-    # Unpickle dataframes
-    if loaded_state['df_original'] is not None:
-        loaded_state['df_original'] = pd.read_pickle(io.BytesIO(loaded_state['df_original']))
-    if loaded_state['df_modified'] is not None:
-        loaded_state['df_modified'] = pd.read_pickle(io.BytesIO(loaded_state['df_modified']))
-        
-    return loaded_state
+        return pickle.load(f)
 
 def prime_data(file_obj, project_name):
     """Main function to load a new CSV, analyze it, and set the initial state."""
     if not file_obj: return init_state()
     df = pd.read_csv(file_obj.name)
     
-    # Smart type conversion
     for col in df.select_dtypes(include=['object']).columns:
         try:
             df[col] = pd.to_datetime(df[col], errors='raise')
@@ -96,35 +96,29 @@ def prime_data(file_obj, project_name):
     
     return {
         "project_name": project_name or f"Project_{datetime.now().strftime('%Y%m%d_%H%M')}",
-        "df_original": df,
-        "df_modified": df.copy(),
-        "metadata": metadata,
-        "insights": insights,
-        "chat_history": [],
-        "dynamic_dashboards": {}
+        "df_original": df, "df_modified": df.copy(), "metadata": metadata,
+        "insights": insights, "chat_history": []
     }
 
 def extract_metadata(df):
     """Utility to get schema and column types."""
     return {
-        'shape': df.shape,
-        'columns': df.columns.tolist(),
+        'shape': df.shape, 'columns': df.columns.tolist(),
         'numeric': df.select_dtypes(include=np.number).columns.tolist(),
         'categorical': df.select_dtypes(include=['object', 'category']).columns.tolist(),
         'datetime': df.select_dtypes(include='datetime').columns.tolist(),
         'dtypes': df.dtypes.apply(lambda x: x.name).to_dict()
     }
 
-# --- Helios Overview Engine ---
+# --- Module-Specific Handlers ---
+
 def run_helios_engine(df, metadata):
-    """The proactive analysis engine."""
+    """The proactive analysis engine for the Helios Overview."""
     insights = {}
-    # Missing Data
     missing = df.isnull().sum()
     insights['missing_data'] = missing[missing > 0].sort_values(ascending=False)
-    # High Cardinality
     insights['high_cardinality'] = {c: df[c].nunique() for c in metadata['categorical'] if df[c].nunique() > 50}
-    # Outlier Detection
+    
     outliers = {}
     for col in metadata['numeric']:
         Q1, Q3 = df[col].quantile(0.25), df[col].quantile(0.75)
@@ -132,7 +126,7 @@ def run_helios_engine(df, metadata):
         count = ((df[col] < (Q1 - 1.5 * IQR)) | (df[col] > (Q3 + 1.5 * IQR))).sum()
         if count > 0: outliers[col] = count
     insights['outliers'] = outliers
-    # ML Target Suggestions
+    
     suggestions = []
     for col in metadata['categorical']:
         if df[col].nunique() == 2: suggestions.append(f"{col} (Classification)")
@@ -141,57 +135,21 @@ def run_helios_engine(df, metadata):
     insights['ml_suggestions'] = suggestions
     return insights
 
-# --- Asclepius Data Lab Handlers ---
-def medic_preview_imputation(state, col, num_method, cat_method):
-    if not col or col not in state['df_modified'].columns: return None
-    df_mod = state['df_modified'].copy()
-    
-    if col in state['metadata']['numeric']:
-        if num_method == 'KNN':
-            imputer = KNNImputer(n_neighbors=5)
-            df_mod[col] = imputer.fit_transform(df_mod[[col]])
-        else:
-            value = df_mod[col].mean() if num_method == 'mean' else df_mod[col].median()
-            df_mod[col].fillna(value, inplace=True)
-        
-        fig = go.Figure()
-        fig.add_trace(go.Histogram(x=state['df_original'][col], name='Original', opacity=0.7))
-        fig.add_trace(go.Histogram(x=df_mod[col], name='Imputed', opacity=0.7))
-        fig.update_layout(barmode='overlay', title_text=f"Distribution for '{col}'", legend_title_text='Dataset')
-        return fig
-    
-    elif col in state['metadata']['categorical']:
-        if cat_method == "Create 'Missing' Category":
-            df_mod[col] = df_mod[col].cat.add_categories("Missing").fillna("Missing") if hasattr(df_mod[col], 'cat') else df_mod[col].fillna("Missing")
-        else: # Mode
-            df_mod[col].fillna(df_mod[col].mode()[0], inplace=True)
-        
-        fig = go.Figure()
-        fig.add_trace(go.Bar(x=state['df_original'][col].value_counts().index, y=state['df_original'][col].value_counts().values, name='Original'))
-        fig.add_trace(go.Bar(x=df_mod[col].value_counts().index, y=df_mod[col].value_counts().values, name='Imputed'))
-        return fig
-    return None
-
-# --- Prometheus Launchpad Handlers ---
 def prometheus_run_model(state, target, features, model_name):
+    """Trains and evaluates a model in the Prometheus Launchpad."""
     if not target or not features: return None, None, "Select target and features."
     df = state['df_modified'].copy()
     df.dropna(subset=[target] + features, inplace=True)
     
-    le_map = {}
     for col in [target] + features:
         if df[col].dtype.name in ['category', 'object']:
-            le = LabelEncoder()
-            df[col] = le.fit_transform(df[col])
-            le_map[col] = le
+            df[col] = LabelEncoder().fit_transform(df[col])
             
     X, y = df[features], df[target]
     problem_type = "Classification" if y.nunique() <= 10 else "Regression"
     
-    MODELS = {
-        "Classification": {"Random Forest": RandomForestClassifier, "Logistic Regression": LogisticRegression},
-        "Regression": {"Random Forest": RandomForestRegressor, "Linear Regression": LinearRegression}
-    }
+    MODELS = {"Classification": {"Random Forest": RandomForestClassifier, "Logistic Regression": LogisticRegression},
+              "Regression": {"Random Forest": RandomForestRegressor, "Linear Regression": LinearRegression}}
     if model_name not in MODELS[problem_type]: return None, None, "Invalid model for this problem type."
     
     model = MODELS[problem_type][model_name](random_state=42)
@@ -199,63 +157,95 @@ def prometheus_run_model(state, target, features, model_name):
     if problem_type == "Classification":
         scores = cross_val_score(model, X, y, cv=5, scoring='accuracy')
         report = f"**Cross-Validated Accuracy:** {np.mean(scores):.3f} ± {np.std(scores):.3f}"
-        
         X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
         model.fit(X_train, y_train)
-        
-        # ROC Curve
         y_prob = model.predict_proba(X_test)[:, 1]
         fpr, tpr, _ = roc_curve(y_test, y_prob)
-        roc_auc = auc(fpr, tpr)
-        fig1 = go.Figure(data=go.Scatter(x=fpr, y=tpr, mode='lines', name=f'ROC curve (area = {roc_auc:.2f})'))
-        fig1.add_scatter(x=[0, 1], y=[0, 1], mode='lines', line=dict(dash='dash'), name='Random Chance')
+        fig1 = go.Figure(data=go.Scatter(x=fpr, y=tpr, mode='lines', name=f'ROC (AUC = {auc(fpr, tpr):.2f})'))
+        fig1.add_scatter(x=[0, 1], y=[0, 1], mode='lines', line=dict(dash='dash'), name='Random')
         fig1.update_layout(title="ROC Curve")
-        
-        # Feature Importance
-        if hasattr(model, 'feature_importances_'):
-            fi = pd.Series(model.feature_importances_, index=features).sort_values(ascending=False)
-            fig2 = px.bar(fi, title="Feature Importance")
-        else: fig2 = go.Figure().update_layout(title="Feature Importance (Not available for this model)")
-            
-        return fig1, fig2, report
     else: # Regression
         scores = cross_val_score(model, X, y, cv=5, scoring='r2')
         report = f"**Cross-Validated R² Score:** {np.mean(scores):.3f} ± {np.std(scores):.3f}"
-        
         X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
         model.fit(X_train, y_train)
         preds = model.predict(X_test)
-        
-        # Residuals Plot
         residuals = y_test - preds
-        fig1 = px.scatter(x=preds, y=residuals, title="Residuals vs. Predicted Plot", labels={'x': 'Predicted Values', 'y': 'Residuals'})
+        fig1 = px.scatter(x=preds, y=residuals, title="Residuals vs. Predicted", labels={'x': 'Predicted', 'y': 'Residuals'})
         fig1.add_hline(y=0, line_dash="dash")
         
-        # Feature Importance
-        if hasattr(model, 'feature_importances_'):
-            fi = pd.Series(model.feature_importances_, index=features).sort_values(ascending=False)
-            fig2 = px.bar(fi, title="Feature Importance")
-        else: fig2 = go.Figure().update_layout(title="Feature Importance (Not available for this model)")
-
-        return fig1, fig2, report
+    if hasattr(model, 'feature_importances_'):
+        fi = pd.Series(model.feature_importances_, index=features).sort_values(ascending=False)
+        fig2 = px.bar(fi, title="Feature Importance")
+    else:
+        fig2 = go.Figure().update_layout(title="Feature Importance (Not available)")
+        
+    return fig1, fig2, report
 
-# --- Athena Co-pilot Handlers ---
 def athena_respond(user_message, history, state, api_key):
-    # Main co-pilot logic
-    pass # This would contain the full logic from previous examples
+    """Handles the chat interaction with the AI Co-pilot."""
+    if not genai:
+        history.append((user_message, "Google AI library not installed. Cannot use Athena."))
+        return history, None, None, state
+    if not api_key:
+        history.append((user_message, "Please enter your Gemini API key to use Athena."))
+        return history, None, None, state
+
+    history.append((user_message, None))
+    
+    # Configure the API
+    genai.configure(api_key=api_key)
+    model = genai.GenerativeModel('gemini-1.5-flash')
 
-def render_dynamic_dashboard(state, dashboard_id):
-    """Renders a dynamically generated dashboard from the state."""
-    # This is a placeholder for the advanced dashboard rendering logic.
-    # In a real scenario, this would execute the Gradio code string stored in state.
-    if dashboard_id in state['dynamic_dashboards']:
-        # This is where we would dynamically create the Gradio components
-        # For this example, we'll return a placeholder
-        return gr.Markdown(f"### Dashboard: {dashboard_id}\n(Dynamic rendering placeholder)")
-    return gr.Markdown("Dashboard not found.")
+    prompt = f"""
+    You are 'Athena', an AI data scientist. Your goal is to help a user by writing and executing Python code on a pandas DataFrame named `df`.
+    
+    **DataFrame Info:**
+    {state['df_modified'].info(verbose=False)}
+    
+    **Instructions:**
+    1.  Analyze the user's request: '{user_message}'.
+    2.  Formulate a plan (thought).
+    3.  Write Python code to execute the plan. You can use `pandas as pd`, `numpy as np`, and `plotly.express as px`.
+    4.  To show a plot, assign it to a variable `fig`.
+    5.  To show a dataframe, assign it to a variable `df_result`.
+    6.  Use `print()` for text output.
+    7.  **NEVER** modify `df` in place.
+    8.  Respond **ONLY** with a single, valid JSON object with keys "thought" and "code".
+    
+    **Your JSON Response:**
+    """
+    try:
+        response = model.generate_content(prompt)
+        response_json = json.loads(response.text.strip().replace("```json", "").replace("```", ""))
+        thought = response_json.get("thought", "Thinking...")
+        code_to_run = response_json.get("code", "print('No code generated.')")
+        
+        bot_thinking = f"🧠 **Thinking:** *{thought}*"
+        history[-1] = (user_message, bot_thinking)
+        yield history, None, None, state
+        
+        local_vars = {'df': state['df_modified'], 'px': px, 'pd': pd, 'np': np}
+        stdout, fig_result, df_result, error = safe_exec(code_to_run, local_vars)
+
+        bot_response = bot_thinking + "\n\n---\n\n"
+        if error: bot_response += f"💥 **Error:**\n```\n{error}\n```"
+        if stdout: bot_response += f"📋 **Output:**\n```\n{stdout}\n```"
+        if not error and not stdout and not fig_result and not isinstance(df_result, pd.DataFrame):
+            bot_response += "✅ Code executed, but produced no direct output."
+
+        history[-1] = (user_message, bot_response)
+        state['chat_history'] = history # Persist chat history
+        yield history, fig_result, df_result, state
+        
+    except Exception as e:
+        error_msg = f"A critical error occurred with the AI model: {e}"
+        history[-1] = (user_message, error_msg)
+        yield history, None, None, state
 
-# --- UI Builder Functions ---
+# --- UI Builder ---
 def build_ui():
+    """Constructs the entire Gradio application interface."""
     with gr.Blocks(theme=THEME, title="Odyssey AI Data Workspace") as demo:
         state = gr.State(init_state())
 
@@ -263,10 +253,10 @@ def build_ui():
             # Left Sidebar - Command Center
             with gr.Column(scale=1):
                 gr.Markdown("# 🦉 Odyssey")
-                
                 with gr.Accordion("📂 Project", open=True):
                     project_name_input = gr.Textbox(label="Project Name", value="New_Project")
                     file_input = gr.File(label="Upload CSV", file_types=[".csv"])
+                    api_key_input = gr.Textbox(label="🔑 Gemini API Key", type="password", placeholder="Enter key...")
                     with gr.Row():
                         save_btn = gr.Button("Save")
                         load_btn = gr.UploadButton("Load .odyssey")
@@ -274,44 +264,24 @@ def build_ui():
                 
                 # Navigation buttons
                 overview_btn = gr.Button(f"{ICONS['overview']} Helios Overview")
-                medic_btn = gr.Button(f"{ICONS['medic']} Asclepius Data Lab")
                 launchpad_btn = gr.Button(f"{ICONS['launchpad']} Prometheus Launchpad")
                 copilot_btn = gr.Button(f"{ICONS['copilot']} Athena Co-pilot")
-                export_btn = gr.Button(f"{ICONS['export']} Export Report")
-
-                # Global Info
-                with gr.Accordion("Global Info", open=False):
-                    file_info_md = gr.Markdown("No file loaded.")
+                export_btn = gr.Button(f"{ICONS['export']} Export Report", visible=False)
 
             # Right Panel - Main Workspace
             with gr.Column(scale=4):
                 # --- Helios Overview Panel ---
                 with gr.Column(visible=True) as overview_panel:
                     gr.Markdown(f"# {ICONS['overview']} Helios Overview")
-                    gr.Markdown("A proactive, high-level summary of your dataset.")
-                    # Interactive dashboard components would go here
-                    helios_report_md = gr.Markdown("Upload data to begin analysis.")
+                    helios_report_md = gr.Markdown("Upload a CSV and provide a project name to begin your Odyssey.")
                 
-                # --- Asclepius Data Lab Panel ---
-                with gr.Column(visible=False) as medic_panel:
-                    gr.Markdown(f"# {ICONS['medic']} Asclepius Data Lab")
-                    gr.Markdown("Interactively clean and prepare your data.")
-                    # UI components for Data Medic
-                    medic_col_select = gr.Dropdown(label="Select Column to Clean")
-                    with gr.Row():
-                        medic_num_method = gr.Radio(['mean', 'median', 'KNN'], label="Numeric Imputation", value='mean')
-                        medic_cat_method = gr.Radio(['mode', "Create 'Missing' Category"], label="Categorical Imputation", value='mode')
-                    medic_preview_plot = gr.Plot()
-                    medic_apply_btn = gr.Button("Apply Changes to Session")
-
                 # --- Prometheus Launchpad Panel ---
                 with gr.Column(visible=False) as launchpad_panel:
                     gr.Markdown(f"# {ICONS['launchpad']} Prometheus Launchpad")
-                    gr.Markdown("Train, evaluate, and understand predictive models.")
-                    # UI components for Launchpad
                     with gr.Row():
                         lp_target = gr.Dropdown(label="🎯 Target")
-                        lp_features = gr.Multiselect(label="✨ Features")
+                        # CORRECTED LINE: Use gr.Dropdown with multiselect=True
+                        lp_features = gr.Dropdown(label="✨ Features", multiselect=True)
                         lp_model = gr.Dropdown(choices=["Random Forest", "Logistic Regression", "Linear Regression"], label="🧠 Model")
                     lp_run_btn = gr.Button("🚀 Launch Model Training (with CV)")
                     lp_report_md = gr.Markdown()
@@ -322,54 +292,59 @@ def build_ui():
                 # --- Athena Co-pilot Panel ---
                 with gr.Column(visible=False) as copilot_panel:
                     gr.Markdown(f"# {ICONS['copilot']} Athena Co-pilot")
-                    gr.Markdown("Your collaborative AI data scientist. Ask anything.")
-                    # Chatbot UI
-                    chatbot = gr.Chatbot(height=500)
-                    with gr.Accordion("AI Generated Dashboard", open=False) as dynamic_dash_accordion:
-                        dynamic_dash_output = gr.Group() # Placeholder for dynamic content
-                    chat_input = gr.Textbox(label="Your Request")
+                    chatbot = gr.Chatbot(height=500, label="Chat History")
+                    with gr.Accordion("AI Generated Results", open=True):
+                        copilot_fig_output = gr.Plot()
+                        copilot_df_output = gr.DataFrame(interactive=False)
+                    chat_input = gr.Textbox(label="Your Request", placeholder="e.g., 'What's the correlation between all numeric columns?'")
                     chat_submit = gr.Button("Send", variant="primary")
         
         # --- Event Handling ---
-        
-        # Panel Navigation
-        panels = [overview_panel, medic_panel, launchpad_panel, copilot_panel]
+        panels = [overview_panel, launchpad_panel, copilot_panel]
         def switch_panel(btn_idx):
             return [gr.update(visible=i == btn_idx) for i in range(len(panels))]
         
         overview_btn.click(lambda: switch_panel(0), None, panels)
-        medic_btn.click(lambda: switch_panel(1), None, panels)
-        launchpad_btn.click(lambda: switch_panel(2), None, panels)
-        copilot_btn.click(lambda: switch_panel(3), None, panels)
+        launchpad_btn.click(lambda: switch_panel(1), None, panels)
+        copilot_btn.click(lambda: switch_panel(2), None, panels)
 
-        # File Upload Logic
-        def on_upload(state, file, name):
-            new_state = prime_data(file, name)
-            # Update all UI components based on the new state
+        def on_upload_or_load(state_data):
+            """Unified function to update UI after data is loaded or a project is loaded."""
             helios_md = "No data loaded."
-            if new_state.get('insights'):
-                helios_md = f"### {ICONS['ml_suggestions']} ML Suggestions\n" + "\n".join([f"- `{s}`" for s in new_state['insights']['ml_suggestions']])
-                # ... Add more sections for a full report
+            if state_data and state_data.get('insights'):
+                insights = state_data['insights']
+                md = f"## 🔭 Proactive Insights for `{state_data.get('project_name')}`\n"
+                md += f"Dataset has **{state_data['metadata']['shape'][0]} rows** and **{state_data['metadata']['shape'][1]} columns**.\n\n"
+                if suggestions := insights.get('ml_suggestions'):
+                    md += "### 🔮 Potential ML Targets\n" + "\n".join(f"- `{s}`" for s in suggestions) + "\n"
+                if not insights.get('missing_data', pd.Series()).empty:
+                    md += "\n### 💧 Missing Data\nFound missing values in these columns:\n" + insights['missing_data'].to_frame('Missing Count').to_markdown() + "\n"
+                helios_md = md
             
-            file_info = f"**File:** `{os.path.basename(file.name)}`\n\n**Shape:** `{new_state['metadata']['shape']}`"
-            
-            all_cols = new_state['metadata']['columns']
-            missing_cols = new_state['insights']['missing_data'].index.tolist()
-
-            return new_state, helios_md, file_info, gr.update(choices=missing_cols), gr.update(choices=all_cols), gr.update(choices=all_cols)
+            all_cols = state_data.get('metadata', {}).get('columns', [])
+            return {
+                state: state_data,
+                helios_report_md: helios_md,
+                lp_target: gr.update(choices=all_cols),
+                lp_features: gr.update(choices=all_cols),
+                chatbot: state_data.get('chat_history', [])
+            }
 
-        file_input.upload(on_upload, [state, file_input, project_name_input], [state, helios_report_md, file_info_md, medic_col_select, lp_target, lp_features])
-        
-        # Project Management
+        file_input.upload(prime_data, [file_input, project_name_input], state).then(
+            on_upload_or_load, state, [state, helios_report_md, lp_target, lp_features, chatbot]
+        )
+        load_btn.upload(load_project, load_btn, state).then(
+            on_upload_or_load, state, [state, helios_report_md, lp_target, lp_features, chatbot]
+        )
         save_btn.click(save_project, state, project_status)
         
-        # Asclepius Live Preview
-        medic_col_select.change(medic_preview_imputation, [state, medic_col_select, medic_num_method, medic_cat_method], medic_preview_plot)
-        medic_num_method.change(medic_preview_imputation, [state, medic_col_select, medic_num_method, medic_cat_method], medic_preview_plot)
-        medic_cat_method.change(medic_preview_imputation, [state, medic_col_select, medic_num_method, medic_cat_method], medic_preview_plot)
-
-        # Prometheus Model Training
         lp_run_btn.click(prometheus_run_model, [state, lp_target, lp_features, lp_model], [lp_fig1, lp_fig2, lp_report_md])
+        
+        chat_submit.click(
+            athena_respond,
+            [chat_input, chatbot, state, api_key_input],
+            [chatbot, copilot_fig_output, copilot_df_output, state]
+        ).then(lambda: "", outputs=chat_input)
 
         return demo