Update app.py

app.py

@@ -332,13 +332,10 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
                     action_col_visible_update, insights_chatbot_visible_update, chatbot_content_update,
                     insights_chat_input_visible_update, insights_suggestions_row_visible_update, 
                     s1_upd, s2_upd, s3_upd, formula_display_visible_update, formula_content_update,
-                    formula_close_hint_md, # This is the component for the hint's visibility, should be formula_close_hint_visible_update
+                    formula_close_hint_visible_update, # Corrected from formula_close_hint_md
                     new_active_action_state_to_set, new_current_chat_plot_id, updated_chat_histories,
                     new_explored_plot_id_to_set 
                 ]
-                # Correcting the hint update:
-                final_updates[10] = formula_close_hint_visible_update
-
 
                 final_updates.extend(generated_panel_vis_updates)   
                 final_updates.extend(generated_bomb_btn_updates)    
@@ -654,113 +651,98 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
             eb_agent_instance_dict = {"agent": None} # To store agent instance across calls if needed, or re-init
 
             async def handle_eb_agent_chat(user_message: str, chat_history_list: list, current_token_state: dict):
+                # Expected outputs: [eb_agent_chatbot_ui, eb_agent_chat_history_st, eb_agent_chat_input_ui, eb_agent_status_md, eb_agent_schema_display_md]
+                # (5 components)
+
                 if not EB_AGENT_AVAILABLE or not os.getenv('GEMINI_API_KEY'):
                     no_key_msg = "L'Agente AI non è disponibile. Assicurati che GEMINI_API_KEY sia configurata."
                     chat_history_list.append([user_message, no_key_msg])
-                    return chat_history_list, chat_history_list, "", gr.update(value=no_key_msg)
+                    # Yield updates for all 5 components
+                    yield chat_history_list, chat_history_list, gr.update(value=""), gr.update(value=no_key_msg), gr.update(value="Nessuno schema disponibile.")
+                    return
 
                 if not user_message.strip():
-                    return chat_history_list, chat_history_list, "", gr.update(value="Stato Agente: Per favore, inserisci una domanda.")
+                    # Yield updates for all 5 components
+                    yield chat_history_list, chat_history_list, gr.update(value=""), gr.update(value="Stato Agente: Per favore, inserisci una domanda."), gr.update() # No change to schema display
+                    return
 
                 status_update_msg = "Stato Agente: Elaborazione della tua richiesta..."
-                yield chat_history_list + [[user_message, None]], chat_history_list + [[user_message, None]], "", gr.update(value=status_update_msg) # Show user message immediately
+                # Show user message immediately, update status
+                yield chat_history_list + [[user_message, None]], chat_history_list + [[user_message, None]], gr.update(value=""), gr.update(value=status_update_msg), gr.update()
 
                 # Prepare DataFrames for the agent
                 df_follower_stats = current_token_state.get("bubble_follower_stats_df", pd.DataFrame())
-                df_posts = current_token_state.get("bubble_posts_df", pd.DataFrame()) # Contains post content
-                df_post_stats = current_token_state.get("bubble_post_stats_df", pd.DataFrame()) # Contains post metrics
+                df_posts = current_token_state.get("bubble_posts_df", pd.DataFrame()) 
+                df_post_stats = current_token_state.get("bubble_post_stats_df", pd.DataFrame()) 
                 df_mentions = current_token_state.get("bubble_mentions_df", pd.DataFrame())
-
-                # For simplicity, we can merge posts and post_stats if they have a common key (e.g., post_id)
-                # Assuming 'id' in df_posts corresponds to 'post_id' in df_post_stats
-                # This is a guess, adjust based on your actual schemas.
-                # If no common key, pass them separately and instruct LLM.
-                df_posts_combined = df_posts
-                if not df_posts.empty and not df_post_stats.empty:
-                    # Attempt merge, ensure column names are correct for your DFs
-                    # Example: common_col_posts = 'id'; common_col_stats = 'post_id'
-                    # For now, let's assume they might not be easily mergeable or LLM can handle separately
-                    pass # LLM will see df_posts and df_post_stats separately
-
-
+                
                 dataframes_for_agent = {
-                    "follower_stats": df_follower_stats.copy() if not df_follower_stats.empty else pd.DataFrame(columns=['no_data_follower_stats']), # Use copy to avoid modifying state
+                    "follower_stats": df_follower_stats.copy() if not df_follower_stats.empty else pd.DataFrame(columns=['no_data_follower_stats']),
                     "posts": df_posts.copy() if not df_posts.empty else pd.DataFrame(columns=['no_data_posts']),
                     "post_stats": df_post_stats.copy() if not df_post_stats.empty else pd.DataFrame(columns=['no_data_post_stats']),
                     "mentions": df_mentions.copy() if not df_mentions.empty else pd.DataFrame(columns=['no_data_mentions'])
                 }
                 
-                # Display schemas for user reference
-                schemas_text = "Schemi DataFrames inviati all'Agente:\n\n"
-                from eb_agent_module import get_all_schemas_representation # Re-import for safety if run in different context
-                schemas_text += get_all_schemas_representation(dataframes_for_agent)
-                # Limit schema display length for UI
+                schemas_text_for_display = "Schemi DataFrames inviati all'Agente:\n\n"
+                from eb_agent_module import get_all_schemas_representation 
+                schemas_text_for_display += get_all_schemas_representation(dataframes_for_agent)
                 max_schema_display_len = 1500
-                if len(schemas_text) > max_schema_display_len:
-                    schemas_text = schemas_text[:max_schema_display_len] + "\n...(schemi troncati per la visualizzazione)"
+                if len(schemas_text_for_display) > max_schema_display_len:
+                    schemas_text_for_display = schemas_text_for_display[:max_schema_display_len] + "\n...(schemi troncati per la visualizzazione)"
                 
-                # Initialize or update agent
-                # For simplicity, re-initialize. If RAG embeddings are slow, consider stateful agent.
                 current_agent = EmployerBrandingAgent(
                     llm_model_name=EB_AGENT_LLM_MODEL,
                     generation_config_params=EB_AGENT_GEN_CONFIG,
                     safety_settings=EB_AGENT_SAFETY_SETTINGS,
                     all_dataframes=dataframes_for_agent,
-                    rag_documents_df=eb_agent_default_rag_docs.copy(), # Use a copy of default RAG docs
+                    rag_documents_df=eb_agent_default_rag_docs.copy(),
                     embedding_model_name=EB_AGENT_EMBEDDING_MODEL,
-                    force_sandbox=True # True to get Python code
+                    force_sandbox=True 
                 )
-                # Restore chat history to the agent if it were stateful
-                # current_agent.chat_history = [msg for pair in chat_history_list for msg_type, msg_content in (('user', pair[0]), ('assistant', pair[1])) if msg_content is not None]
-                # For stateless re-init, we pass history to process_query if needed by LLM, but agent itself starts fresh or manages its own history via its methods
-
-                # Simplified history for this agent: it manages its own via process_query
-                # We only need to pass the user_message.
-                # The agent's process_query will handle its internal history.
                 
-                # If the agent's internal history needs to match the UI history:
-                # Convert Gradio chat history to agent's expected format
                 agent_internal_history = []
-                for user_q, ai_r in chat_history_list:
+                for user_q, ai_r in chat_history_list: # Use the passed chat_history_list
                     if user_q: agent_internal_history.append({"role": "user", "content": user_q})
-                    if ai_r: agent_internal_history.append({"role": "assistant", "content": ai_r}) # or "model"
-                
-                # Update agent's history before processing (if agent supports this way)
-                current_agent.chat_history = agent_internal_history # Agent will append new query and response
+                    if ai_r: agent_internal_history.append({"role": "assistant", "content": ai_r})
+                current_agent.chat_history = agent_internal_history
 
                 try:
                     logging.info(f"Sending to EB Agent. User: '{user_message}'. DF Keys: {list(dataframes_for_agent.keys())}")
-                    # The agent's process_query already appends to its internal history
                     ai_response = await current_agent.process_query(user_query=user_message)
                     
-                    # The agent's internal history is now updated. We need to reflect this in Gradio's history state.
-                    # The last two entries in agent.chat_history are the current user_message and ai_response
                     updated_gradio_history = []
-                    temp_hist = current_agent.chat_history
+                    temp_hist = current_agent.chat_history 
                     for i in range(0, len(temp_hist), 2):
                         u_msg = temp_hist[i]['content']
-                        a_msg = temp_hist[i+1]['content'] if i+1 < len(temp_hist) else "Thinking..." # Should have assistant response
+                        a_msg = temp_hist[i+1]['content'] if i+1 < len(temp_hist) else "Thinking..." 
                         updated_gradio_history.append([u_msg, a_msg])
                     
                     status_update_msg = "Stato Agente: Risposta ricevuta."
-                    yield updated_gradio_history, updated_gradio_history, "", gr.update(value=status_update_msg), gr.update(value=schemas_text)
+                    # Yield final updates for all 5 components
+                    yield updated_gradio_history, updated_gradio_history, gr.update(value=""), gr.update(value=status_update_msg), gr.update(value=schemas_text_for_display)
 
                 except Exception as e:
                     logging.error(f"Error during EB Agent processing: {e}", exc_info=True)
                     error_msg = f"# Errore dell'Agente AI:\n{type(e).__name__}: {str(e)}"
-                    # Use the last known good history before error for UI state
-                    chat_history_list.append([user_message, error_msg]) # Add error to current turn
+                    
+                    # Ensure the current turn in chat_history_list reflects the error
+                    # The last item in chat_history_list is [user_message, None] from the first yield
+                    current_turn_history = chat_history_list + [[user_message, error_msg]] # This might duplicate user message if not careful
+                    
+                    # Let's reconstruct history carefully to avoid duplicates
+                    final_error_history = chat_history_list # This already has [user_message, None] as last item if first yield happened
+                    if final_error_history and final_error_history[-1][0] == user_message and final_error_history[-1][1] is None:
+                        final_error_history[-1][1] = error_msg # Update the assistant part of the last entry
+                    else: # Should not happen if first yield was correct
+                        final_error_history.append([user_message, error_msg])
+
+
                     status_update_msg = f"Stato Agente: Errore - {type(e).__name__}"
-                    yield chat_history_list, chat_history_list, "", gr.update(value=status_update_msg), gr.update(value=schemas_text)
+                    # Yield error updates for all 5 components
+                    yield final_error_history, final_error_history, gr.update(value=""), gr.update(value=status_update_msg), gr.update(value=schemas_text_for_display)
 
 
             def clear_eb_agent_chat_history():
-                # eb_agent_instance_dict["agent"] might need its history cleared too if stateful
-                # For re-init agent, just clearing UI state is enough.
-                # If agent instance is stored and reused:
-                # agent_to_clear = eb_agent_instance_dict.get("agent")
-                # if agent_to_clear and hasattr(agent_to_clear, 'clear_chat_history'):
-                #    agent_to_clear.clear_chat_history()
                 initial_msg = "Ciao! Sono il tuo Agente AI per l'Employer Branding. Come posso aiutarti?" if EB_AGENT_AVAILABLE else "Agente AI non disponibile."
                 return [[None, initial_msg]], [[None, initial_msg]], "Stato Agente: Chat resettata."