upgrade memoria + UX

app.py

@@ -435,13 +435,17 @@ def convert_to_messages_format(chat_history):
     
     return messages
 
-async def stream_agent_response(question: str, chat_history: List[List[str]]) -> Tuple[str, Optional["go.Figure"]]:
-    """Procesa la pregunta del usuario y devuelve la respuesta del agente con memoria de conversación."""
+async def stream_agent_response(question: str, chat_history: List[List[str]]) -> Tuple[str, Optional["go.Figure"], Optional[Dict[str, Any]]]:
+    """Procesa la pregunta del usuario y devuelve (respuesta, figura, chart_state).
+
+    chart_state schema: {"data": List[Dict], "x_col": str, "y_col": str, "title": str, "chart_type": str}
+    """
     global agent  # Make sure we can modify the agent's memory
     
     # Initialize response
     response_text = ""
     chart_fig = None
+    chart_state: Optional[Dict[str, Any]] = None
     messages = []
     
     # Add previous chat history in the correct format for the agent
@@ -508,7 +512,6 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
             sql_query = extract_sql_query(response_text)
             if sql_query and looks_like_sql(sql_query):
                 logger.info(f"Detected SQL query: {sql_query}")
-                # Execute the query and update the response
                 db_connection, _ = setup_database_connection()
                 if db_connection:
                     query_result = execute_sql_query(sql_query, db_connection)
@@ -550,24 +553,26 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
 
                                     # Choose x/y columns (assume first is category, second numeric)
                                     x_col = columns[0]
-                                    y_col = columns[1]
-
-                                    # Coerce numeric values for y
-                                    for row in data:
+                                    # pick first numeric column different to x
+                                    y_col = None
+                                    for col in columns[1:]:
                                         try:
-                                            row[y_col] = float(re.sub(r"[^0-9.\-]", "", str(row[y_col])))
+                                            pd.to_numeric(data[0][col])
+                                            y_col = col
+                                            break
                                         except Exception:
-                                            pass
-
-                                    chart_fig = generate_chart(
-                                        data=data,
-                                        chart_type=desired_type,
-                                        x=x_col,
-                                        y=y_col,
-                                        title=f"{y_col} por {x_col}"
-                                    )
-                                    if chart_fig is not None:
-                                        logger.info(f"Chart generated from SQL table: type={desired_type}, x={x_col}, y={y_col}, rows={len(data)}")
+                                            continue
+                                    if y_col:
+                                        chart_fig = generate_chart(
+                                            data=data,
+                                            chart_type=desired_type,
+                                            x=x_col,
+                                            y=y_col,
+                                            title=f"{y_col} por {x_col}"
+                                        )
+                                        if chart_fig is not None:
+                                            logger.info(f"Chart generated from SQL table: type={desired_type}, x={x_col}, y={y_col}, rows={len(data)}")
+                                            chart_state = {"data": data, "x_col": x_col, "y_col": y_col, "title": f"{y_col} por {x_col}", "chart_type": desired_type}
                     except Exception as e:
                         logger.error(f"Error generating chart: {str(e)}", exc_info=True)
                         # Don't fail the whole request if chart generation fails
@@ -630,6 +635,7 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
                                         )
                                         if chart_fig is not None:
                                             logger.info("Chart generated from second-pass SQL execution.")
+                                            chart_state = {"data": data, "x_col": x_col, "y_col": y_col, "title": f"{y_col} por {x_col}", "chart_type": desired_type}
                             else:
                                 logger.info("No DB connection on second pass; skipping.")
                     except Exception as e:
@@ -687,6 +693,7 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
                             )
                             if chart_fig is not None:
                                 logger.info(f"Chart generated from text fallback: type={desired_type}, items={len(data)}")
+                                chart_state = {"data": data, "x_col": "label", "y_col": "value", "title": "Distribución", "chart_type": desired_type}
             
             # Update the assistant's message with the response
             assistant_message["content"] = response_text
@@ -713,7 +720,7 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
             logger.info("No chart generated for this turn.")
         else:
             logger.info("Returning a chart figure to UI.")
-        return message_content, chart_fig
+        return message_content, chart_fig, chart_state
         
     except Exception as e:
         error_msg = f"## ❌ Error\n\nOcurrió un error al procesar tu solicitud:\n\n```\n{str(e)}\n```"
@@ -924,13 +931,16 @@ def create_ui():
         
         # Hidden component for streaming output
         streaming_output_display = gr.Textbox(visible=False)
+
+        # Session state to persist last chart data for follow-ups
+        chart_state = gr.State(value=None)
         
-        return demo, chatbot, chart_display, question_input, submit_button, streaming_output_display
+        return demo, chatbot, chart_display, question_input, submit_button, streaming_output_display, chart_state
 
 def create_application():
     """Create and configure the Gradio application."""
     # Create the UI components
-    demo, chatbot, chart_display, question_input, submit_button, streaming_output_display = create_ui()
+    demo, chatbot, chart_display, question_input, submit_button, streaming_output_display, chart_state = create_ui()
     
     def user_message(user_input: str, chat_history: List[Dict[str, str]]) -> Tuple[str, List[Dict[str, str]]]:
         """Add user message to chat history (messages format) and clear input."""
@@ -947,10 +957,14 @@ def create_application():
 
         return "", chat_history
     
-    async def bot_response(chat_history: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Optional[go.Figure]]:
-        """Generate bot response for messages-format chat history and return optional chart figure."""
+    async def bot_response(chat_history: List[Dict[str, str]], chart_state_value: Optional[Dict[str, Any]]) -> Tuple[List[Dict[str, str]], Optional[go.Figure], Optional[Dict[str, Any]]]:
+        """Generate bot response and optional chart figure using messages-format chat history.
+
+        Also accepts and returns chart_state (dict) to support follow-up prompts like
+        'muestra los mismos datos en barras'.
+        """
         if not chat_history:
-            return chat_history, None
+            return chat_history, None, chart_state_value
 
         # Ensure last message is a user turn awaiting assistant reply
         last = chat_history[-1]
@@ -961,6 +975,44 @@ def create_application():
             question = last["content"]
             logger.info(f"Processing question: {question}")
 
+            # Detect quick follow-up re-plot using the same data
+            ql = question.lower()
+            wants_reuse = any(kw in ql for kw in [
+                "mismos datos", "estos mismos", "los mismos", "igual data", "misma data", "con los mismos"
+            ])
+            # Determine requested chart type if any
+            desired_type = None
+            if any(k in ql for k in ["barras", "bar", "columnas"]):
+                desired_type = "bar"
+            elif any(k in ql for k in ["línea", "linea", "line"]):
+                desired_type = "line"
+            elif any(k in ql for k in ["pastel", "pie", "circular"]):
+                desired_type = "pie"
+            elif "scatter" in ql or "dispersión" in ql or "dispersion" in ql:
+                desired_type = "scatter"
+            elif "histograma" in ql or "histogram" in ql:
+                desired_type = "histogram"
+
+            if wants_reuse and chart_state_value and isinstance(chart_state_value, dict) and chart_state_value.get("data"):
+                # Re-plot with same data using requested type (or keep previous if not specified)
+                new_type = desired_type or chart_state_value.get("chart_type", "bar")
+                fig = generate_chart(
+                    data=chart_state_value.get("data"),
+                    chart_type=new_type,
+                    x=chart_state_value.get("x_col", "label"),
+                    y=chart_state_value.get("y_col", "value"),
+                    title=chart_state_value.get("title", "Distribución")
+                )
+                if fig is not None:
+                    friendly = f"He actualizado la visualización a {('gráfico de ' + new_type) if new_type != 'pie' else 'gráfico circular'} usando los mismos datos."
+                    chat_history.append({"role": "assistant", "content": friendly})
+                    # Save new type in state
+                    chart_state_value = {
+                        **chart_state_value,
+                        "chart_type": new_type,
+                    }
+                    return chat_history, fig, chart_state_value
+
             # Convert prior messages to pair history for stream_agent_response()
             pair_history: List[List[str]] = []
             i = 0
@@ -979,20 +1031,23 @@ def create_application():
                     i += 1
 
             # Call the agent for this new user question
-            assistant_message, chart_fig = await stream_agent_response(question, pair_history)
+            assistant_message, chart_fig, new_state = await stream_agent_response(question, pair_history)
 
             # Append assistant message back into messages history
             chat_history.append({"role": "assistant", "content": assistant_message})
 
+            # Update chart_state if provided by the agent path
+            chart_state_value = new_state if new_state is not None else chart_state_value
+
             logger.info("Response generation complete")
-            return chat_history, chart_fig
+            return chat_history, chart_fig, chart_state_value
 
         except Exception as e:
             error_msg = f"## ❌ Error\n\nError al procesar la solicitud:\n\n```\n{str(e)}\n```"
             logger.error(error_msg, exc_info=True)
             # Ensure we add an assistant error message for the UI
             chat_history.append({"role": "assistant", "content": error_msg})
-            return chat_history, None
+            return chat_history, None, chart_state_value
     
     # Event handlers
     with demo:
@@ -1004,8 +1059,8 @@ def create_application():
             queue=True
         ).then(
             fn=bot_response,
-            inputs=[chatbot],
-            outputs=[chatbot, chart_display],
+            inputs=[chatbot, chart_state],
+            outputs=[chatbot, chart_display, chart_state],
             api_name="ask"
         )
         
@@ -1017,8 +1072,8 @@ def create_application():
             queue=True
         ).then(
             fn=bot_response,
-            inputs=[chatbot],
-            outputs=[chatbot, chart_display]
+            inputs=[chatbot, chart_state],
+            outputs=[chatbot, chart_display, chart_state]
         )
     
     return demo