Update app.py

app.py

@@ -652,28 +652,31 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
             
             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])
-                    # 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.")
+                    # Ensure chat_history_list is mutable if it comes from gr.State
+                    current_chat_history = list(chat_history_list) if chat_history_list else []
+                    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 []
+            
                 if not user_message.strip():
-                    # 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
+                    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
             
                 status_update_msg = "Stato Agente: Elaborazione della tua richiesta..."
                 # 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()
+                # Add user message to current history before yielding
+                pending_history = current_chat_history + [[user_message, None]]
+                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()) 
-                df_post_stats = current_token_state.get("bubble_post_stats_df", pd.DataFrame()) 
+                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())
                 
                 dataframes_for_agent = {
@@ -684,13 +687,13 @@ 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 
-                schemas_text_for_display += get_all_schemas_representation(dataframes_for_agent)
+                # 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
                 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(
+                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,
@@ -699,44 +702,96 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
                     embedding_model_name=EB_AGENT_EMBEDDING_MODEL
                 )
                 
-                # Initialize the agent's chat history with previous conversation
                 agent_internal_history = []
-                for user_q, ai_r in chat_history_list:
+                for user_q, ai_r_obj in current_chat_history: # Iterate over the current history being built
                     if user_q: agent_internal_history.append({"role": "user", "content": user_q})
-                    if ai_r: agent_internal_history.append({"role": "model", "content": ai_r})
-                
-                # ADD THE CURRENT USER MESSAGE TO THE AGENT'S HISTORY - THIS WAS MISSING!
+                    # ai_r_obj could be string, tuple (text, image_url), or None
+                    if ai_r_obj:
+                        if isinstance(ai_r_obj, tuple):
+                            # If it's a (text, image_url) tuple, take the text part for agent's history
+                            # Or combine them if your agent can handle it. For simplicity, just text.
+                            text_for_agent_history = ai_r_obj[0] if ai_r_obj[0] else "Visual media displayed."
+                            agent_internal_history.append({"role": "model", "content": text_for_agent_history})
+                        elif isinstance(ai_r_obj, str):
+                            agent_internal_history.append({"role": "model", "content": ai_r_obj})
+            
+                # ADD THE CURRENT USER MESSAGE TO THE AGENT'S HISTORY
                 agent_internal_history.append({"role": "user", "content": user_message})
                 current_agent.chat_history = agent_internal_history
-            
+                
                 try:
-                    # Initialize the agent first
                     init_success = await current_agent.initialize()
                     if not init_success:
                         error_msg = "Errore: Impossibile inizializzare l'agente AI."
-                        updated_history = chat_history_list + [[user_message, error_msg]]
+                        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
             
                     logging.info(f"Sending to EB Agent. User: '{user_message}'. DF Keys: {list(dataframes_for_agent.keys())}")
-                    ai_response = await current_agent.process_query(user_query=user_message)
+                    # 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) 
                     
-                    # Update the chat history with the AI response
-                    updated_history = chat_history_list + [[user_message, ai_response]]
+                    bot_message_for_display = "Error: Agent returned an unexpected response." # Default
+            
+                    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
+                        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}")
+                    
+                    elif isinstance(ai_response_dict, str): # Agent returned a plain string
+                        bot_message_for_display = ai_response_dict
+                    else: # Fallback for other unexpected types
+                        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]]
                     
                     status_update_msg = "Stato Agente: Risposta ricevuta."
-                    # Yield final updates for all 5 components
                     yield 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 = f"# Errore dell'Agente AI:\n{type(e).__name__}: {str(e)}"
-                    
-                    # Update history with error message
-                    updated_history = chat_history_list + [[user_message, error_msg]]
-            
+                    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]]
                     status_update_msg = f"Stato Agente: Errore - {type(e).__name__}"
-                    # Yield error updates for all 5 components
                     yield updated_history, updated_history, gr.update(value=""), gr.update(value=status_update_msg), gr.update(value=schemas_text_for_display)