Update app.py

app.py

@@ -650,28 +650,23 @@ 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]
-                
+                current_chat_history = list(chat_history_list) if chat_history_list else []
+    
                 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."
-                    # Ensure chat_history_list is mutable if it comes from gr.State
-                    current_chat_history = list(chat_history_list) if chat_history_list else []
+                    no_key_msg = "L'Agente AI non è disponibile. Assicurati che GEMINI_API_KEY sia configurata e che l'agente sia caricato."
                     current_chat_history.append([user_message, no_key_msg])
-                    yield current_chat_history, current_chat_history, gr.update(value=""), gr.update(value=no_key_msg), gr.update(value="Nessuno schema disponibile.")
-                    return
-            
-                current_chat_history = list(chat_history_list) if chat_history_list else []
-            
+                    # Ensure all outputs are updated
+                    return current_chat_history, current_chat_history, gr.update(value=""), gr.update(value=no_key_msg), gr.update(value="Nessuno schema disponibile.")
+    
                 if not user_message.strip():
-                    yield current_chat_history, current_chat_history, gr.update(value=""), gr.update(value="Stato Agente: Per favore, inserisci una domanda."), gr.update() # No change to schema display
-                    return
-            
+                    # Ensure all outputs are updated even for empty message
+                    return current_chat_history, current_chat_history, gr.update(value=""), gr.update(value="Stato Agente: Per favore, inserisci una domanda."), gr.update()
+    
                 status_update_msg = "Stato Agente: Elaborazione della tua richiesta..."
-                # Show user message immediately, update status
-                # Add user message to current history before yielding
                 pending_history = current_chat_history + [[user_message, None]]
+                # Yield intermediate state to show user message immediately
                 yield pending_history, pending_history, 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())  
@@ -686,18 +681,20 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
                 }
                 
                 schemas_text_for_display = "Schemi DataFrames inviati all'Agente:\n\n"
-                from eb_agent_module import get_all_schemas_representation # Assuming this is correctly imported in your main file
-                schemas_text_for_display += get_all_schemas_representation(dataframes_for_agent) # Using the mock or your actual function
+                schemas_text_for_display += get_all_schemas_representation(dataframes_for_agent)
                 max_schema_display_len = 1500
                 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)"
                 
-                current_agent = EmployerBrandingAgent( # Using the mock or your actual class
-                    llm_model_name=EB_AGENT_LLM_MODEL,
-                    generation_config_dict=EB_AGENT_GEN_CONFIG,
-                    safety_settings_list=EB_AGENT_SAFETY_SETTINGS,
+                # Update schema display before agent call
+                yield pending_history, pending_history, gr.update(value=""), gr.update(value=status_update_msg), gr.update(value=schemas_text_for_display)
+    
+                current_agent = EmployerBrandingAgent(
+                    llm_model_name="gemini-2.5-flash-preview-05-20", # Replace with your actual EB_AGENT_LLM_MODEL
+                    generation_config_dict={}, # Replace with your actual EB_AGENT_GEN_CONFIG
+                    safety_settings_list=[], # Replace with your actual EB_AGENT_SAFETY_SETTINGS
                     all_dataframes=dataframes_for_agent,
-                    embedding_model_name=EB_AGENT_EMBEDDING_MODEL
+                    embedding_model_name="gemini-embedding-exp-03-07" # Replace with your actual EB_AGENT_EMBEDDING_MODEL
                 )
                 
                 agent_internal_history = []
@@ -721,81 +718,67 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
                     init_success = await current_agent.initialize()
                     if not init_success:
                         error_msg = "Errore: Impossibile inizializzare l'agente AI."
-                        updated_history = current_chat_history + [[user_message, error_msg]]
-                        yield updated_history, updated_history, gr.update(value=""), gr.update(value="Stato Agente: Errore di inizializzazione"), gr.update(value=schemas_text_for_display)
-                        return
-            
+                        # Use pending_history which already has the user message
+                        updated_history = pending_history[:-1] + [[user_message, error_msg]] # Replace None with error
+                        return updated_history, updated_history, gr.update(value=""), gr.update(value="Stato Agente: Errore di inizializzazione"), gr.update(value=schemas_text_for_display)
+                
                     logging.info(f"Sending to EB Agent. User: '{user_message}'. DF Keys: {list(dataframes_for_agent.keys())}")
-                    # ai_response_dict is what the agent returns. Based on error, it's {'text': 'blob...'}
-                    ai_response_dict = await current_agent.process_query(user_query=user_message) 
+                    ai_response_dict = await current_agent.process_query(user_query=user_message)  
                     
                     bot_message_for_display = "Error: Agent returned an unexpected response." # Default
-            
+                    agent_text_response = None
+                    agent_image_path = None
+    
                     if isinstance(ai_response_dict, dict):
-                        combined_message_blob = ai_response_dict.get("text")
-            
-                        if isinstance(combined_message_blob, str):
-                            text_part = combined_message_blob
-                            image_data_url = None
-                            
-                            # Attempt to parse image data URL from the combined_message_blob
-                            # This assumes the image data URL, if present, is on its own line or at the end.
-                            lines = combined_message_blob.splitlines()
-                            if lines:
-                                possible_image_prefixes = [
-                                    "data:image/png;base64,", 
-                                    "data:image/jpeg;base64,", 
-                                    "data:image/gif;base64,",
-                                    "data:image/webp;base64,"
-                                ]
-                                # Check lines from the end, as plot is likely at the end of the message
-                                for i in range(len(lines) - 1, -1, -1):
-                                    current_line = lines[i].strip()
-                                    for prefix in possible_image_prefixes:
-                                        if current_line.startswith(prefix):
-                                            # Basic validation: check for typical base64 characters and some length
-                                            # This is a heuristic to ensure it's likely a valid base64 data string
-                                            if len(current_line) > len(prefix) + 20 and \
-                                               all(c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=" for c in current_line[len(prefix):]):
-                                                image_data_url = current_line
-                                                # Reconstruct text_part from lines *before* this image line
-                                                text_part = "\n".join(lines[:i]).strip()
-                                                break # Found image prefix
-                                    if image_data_url:
-                                        break # Found image line
-            
-                            if image_data_url:
-                                # If text_part became empty after extracting image, use None for text in tuple
-                                bot_message_for_display = (text_part if text_part else None, image_data_url)
-                            else:
-                                # No image found or parsing failed, treat the whole blob as text
-                                bot_message_for_display = combined_message_blob
+                        agent_text_response = ai_response_dict.get("text")
+                        agent_image_path = ai_response_dict.get("image_path")
+    
+                        logging.info(f"Agent response: text='{str(agent_text_response)[:100]}...', image_path='{agent_image_path}'")
+    
+                        if agent_image_path and os.path.exists(agent_image_path):
+                            # If there's an image, display it. Text can accompany it or be None.
+                            # Ensure text_response is a string or None for the tuple.
+                            text_for_display = agent_text_response if isinstance(agent_text_response, str) else None
+                            if not text_for_display and agent_text_response is not None: # e.g. if text_response was empty string
+                                 text_for_display = str(agent_text_response)
+    
+                            bot_message_for_display = (text_for_display, agent_image_path)
+                            logging.info(f"Displaying image: {agent_image_path} with text: {text_for_display}")
+                        elif agent_text_response is not None: # Text only, or image path was invalid
+                            bot_message_for_display = agent_text_response
+                            if agent_image_path: # Log if image path was given but not valid
+                                logging.warning(f"Agent provided image_path '{agent_image_path}' but it does not exist or was not used.")
                         else:
-                            bot_message_for_display = "Agent returned a dictionary, but the 'text' field was not a string or was missing."
-                            logging.warning(f"AI response dict 'text' field issue. Dict: {ai_response_dict}")
+                            bot_message_for_display = "L'agente ha risposto, ma il contenuto non è visualizzabile."
+                            logging.warning(f"AI response dict issue. Text: {agent_text_response}, Image Path: {agent_image_path}")
                     
-                    elif isinstance(ai_response_dict, str): # Agent returned a plain string
+                    elif isinstance(ai_response_dict, str): 
                         bot_message_for_display = ai_response_dict
-                    else: # Fallback for other unexpected types
+                        logging.warning(f"AI response was a plain string: {ai_response_dict}")
+                    else: 
                         bot_message_for_display = f"Error: Agent returned an unexpected data type: {type(ai_response_dict)}."
                         logging.error(f"Unexpected AI response type: {type(ai_response_dict)}, content: {ai_response_dict}")
-            
-                    updated_history = current_chat_history + [[user_message, bot_message_for_display]]
+                    
+                    # Use pending_history which already has the user message with None for bot response
+                    updated_history = pending_history[:-1] + [[user_message, bot_message_for_display]]
                     
                     status_update_msg = "Stato Agente: Risposta ricevuta."
-                    yield updated_history, updated_history, gr.update(value=""), gr.update(value=status_update_msg), gr.update(value=schemas_text_for_display)
-            
+                    return updated_history, updated_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_for_chat = f"# Errore dell'Agente AI:\n{type(e).__name__}: {str(e)}"
-                    updated_history = current_chat_history + [[user_message, error_msg_for_chat]]
+                    # Use pending_history which already has the user message
+                    updated_history = pending_history[:-1] + [[user_message, error_msg_for_chat]]
                     status_update_msg = f"Stato Agente: Errore - {type(e).__name__}"
-                    yield updated_history, updated_history, gr.update(value=""), gr.update(value=status_update_msg), gr.update(value=schemas_text_for_display)
+                    return updated_history, updated_history, gr.update(value=""), gr.update(value=status_update_msg), gr.update(value=schemas_text_for_display)
 
 
             def clear_eb_agent_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."
+                # Ensure all outputs are updated
+                return [[None, initial_msg]], [[None, initial_msg]], "Stato Agente: Chat resettata.", gr.update(value="Gli schemi dei dati verranno mostrati qui...")
+
 
             # Connect UI to Handler for EB Agent
             eb_agent_submit_btn.click(