upgrade fix chatgpt5

app.py

@@ -43,21 +43,21 @@ def generate_chart(data: Union[Dict, List[Dict], pd.DataFrame],
                  y: str = None, 
                  title: str = "", 
                  x_label: str = None, 
-                 y_label: str = None) -> str:
+                 y_label: str = None):
     """
-    Generate a chart from data and return it as a base64 encoded image.
+    Generate an interactive Plotly figure from data.
     
     Args:
         data: The data to plot (can be a list of dicts or a pandas DataFrame)
         chart_type: Type of chart to generate (bar, line, pie, scatter, histogram)
-        x: Column name for x-axis
-        y: Column name for y-axis (not needed for pie charts)
+        x: Column name for x-axis (names for pie)
+        y: Column name for y-axis (values for pie)
         title: Chart title
         x_label: Label for x-axis
         y_label: Label for y-axis
         
     Returns:
-        Markdown string with embedded image
+        A Plotly Figure object (interactive) or None on error
     """
     try:
         # Convert data to DataFrame if it's a list of dicts
@@ -69,7 +69,7 @@ def generate_chart(data: Union[Dict, List[Dict], pd.DataFrame],
             df = data
             
         if not isinstance(df, pd.DataFrame):
-            return "Error: Data must be a dictionary, list of dictionaries, or pandas DataFrame"
+            return None
             
         # Generate the appropriate chart type
         fig = None
@@ -78,13 +78,13 @@ def generate_chart(data: Union[Dict, List[Dict], pd.DataFrame],
         elif chart_type == 'line':
             fig = px.line(df, x=x, y=y, title=title)
         elif chart_type == 'pie':
-            fig = px.pie(df, names=x, values=y, title=title)
+            fig = px.pie(df, names=x, values=y, title=title, hole=0)
         elif chart_type == 'scatter':
             fig = px.scatter(df, x=x, y=y, title=title)
         elif chart_type == 'histogram':
             fig = px.histogram(df, x=x, title=title)
         else:
-            return "Error: Unsupported chart type. Use 'bar', 'line', 'pie', 'scatter', or 'histogram'"
+            return None
         
         # Update layout
         fig.update_layout(
@@ -96,24 +96,13 @@ def generate_chart(data: Union[Dict, List[Dict], pd.DataFrame],
             height=400
         )
         
-        # Save the figure to a temporary file
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
-        fig.write_image(temp_file.name, format='png', engine='kaleido')
-        
-        # Read the image file and encode as base64
-        with open(temp_file.name, 'rb') as img_file:
-            img_base64 = base64.b64encode(img_file.read()).decode('utf-8')
-        
-        # Clean up the temporary file
-        os.unlink(temp_file.name)
-        
-        # Return as markdown image
-        return f'<img src="data:image/png;base64,{img_base64}" style="max-width:100%;"/>'
+        return fig
         
     except Exception as e:
         error_msg = f"Error generating chart: {str(e)}"
         logger.error(error_msg, exc_info=True)
-        return f"<div style='color: red;'>{error_msg}</div>"
+        return None
+
 logger = logging.getLogger(__name__)
 
 def check_environment():
@@ -423,12 +412,13 @@ def convert_to_messages_format(chat_history):
     
     return messages
 
-async def stream_agent_response(question: str, chat_history: List[List[str]]) -> str:
+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."""
     global agent  # Make sure we can modify the agent's memory
     
     # Initialize response
     response_text = ""
+    chart_fig = None
     messages = []
     
     # Add previous chat history in the correct format for the agent
@@ -505,7 +495,7 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
                     # Add the query and its result to the response
                     response_text += f"\n\n### 🔍 Resultado de la consulta:\n```sql\n{sql_query}\n```\n\n{query_result}"
                     
-                    # Try to generate a chart if the result is tabular
+                    # Try to generate an interactive chart if the result is tabular
                     try:
                         if isinstance(query_result, str) and '|' in query_result and '---' in query_result:
                             # Convert markdown table to DataFrame
@@ -526,32 +516,37 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
                                         data.append(dict(zip(columns, values)))
                                 
                                 if data and len(columns) >= 2:
-                                    # Generate a chart based on the data
-                                    chart_type = 'bar'  # Default chart type
-                                    if len(columns) == 2:
-                                        # Simple bar chart for two columns
-                                        chart_html = generate_chart(
-                                            data=data,
-                                            chart_type=chart_type,
-                                            x=columns[0],
-                                            y=columns[1],
-                                            title=f"{columns[1]} por {columns[0]}",
-                                            x_label=columns[0],
-                                            y_label=columns[1]
-                                        )
-                                        response_text += f"\n\n### 📊 Visualización:\n{chart_html}"
-                                    elif len(columns) > 2:
-                                        # For multiple columns, create a line chart
-                                        chart_html = generate_chart(
-                                            data=data,
-                                            chart_type='line',
-                                            x=columns[0],
-                                            y=columns[1],
-                                            title=f"{', '.join(columns[1:])} por {columns[0]}",
-                                            x_label=columns[0],
-                                            y_label=", ".join(columns[1:])
-                                        )
-                                        response_text += f"\n\n### 📊 Visualización:\n{chart_html}"
+                                    # Determine chart type from user's question (supports pie chart)
+                                    q_lower = question.lower()
+                                    if any(k in q_lower for k in ["gráfico circular", "grafico circular", "pie", "pastel"]):
+                                        desired_type = 'pie'
+                                    elif any(k in q_lower for k in ["línea", "linea", "line"]):
+                                        desired_type = 'line'
+                                    elif any(k in q_lower for k in ["dispersión", "dispersion", "scatter"]):
+                                        desired_type = 'scatter'
+                                    elif any(k in q_lower for k in ["histograma", "histogram"]):
+                                        desired_type = 'histogram'
+                                    else:
+                                        desired_type = 'bar'
+
+                                    # 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:
+                                        try:
+                                            row[y_col] = float(re.sub(r"[^0-9.\-]", "", str(row[y_col])))
+                                        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}"
+                                    )
                     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
@@ -579,15 +574,14 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
         else:
             message_content = str(assistant_message)
             
-        # Return the assistant's response in the format expected by the bot_response function
-        # The bot_response function will handle updating the chat history
-        return message_content
+        # Return the assistant's response and an optional interactive chart figure
+        return message_content, chart_fig
         
     except Exception as e:
         error_msg = f"## ❌ Error\n\nOcurrió un error al procesar tu solicitud:\n\n```\n{str(e)}\n```"
         logger.error(f"Error in stream_agent_response: {str(e)}", exc_info=True)
-        # Ensure we return in the correct format: [(user_msg, bot_msg)]
-        return [(None, error_msg)]
+        # Return error message and no chart
+        return error_msg, None
 
 # Custom CSS for the app
 custom_css = """
@@ -703,6 +697,13 @@ def create_ui():
                 layout="panel"  # Better layout for messages
             )
         
+        # Chart display area (interactive Plotly figure)
+        chart_display = gr.Plot(
+            label="📊 Visualización",
+            height=420,
+            interactive=True,
+        )
+        
         # Input area
         with gr.Row():
             question_input = gr.Textbox(
@@ -788,12 +789,12 @@ def create_ui():
         # Hidden component for streaming output
         streaming_output_display = gr.Textbox(visible=False)
         
-        return demo, chatbot, question_input, submit_button, streaming_output_display
+        return demo, chatbot, chart_display, question_input, submit_button, streaming_output_display
 
 def create_application():
     """Create and configure the Gradio application."""
     # Create the UI components
-    demo, chatbot, question_input, submit_button, streaming_output_display = create_ui()
+    demo, chatbot, chart_display, question_input, submit_button, streaming_output_display = create_ui()
     
     def user_message(user_input: str, chat_history: List[List[str]]) -> Tuple[str, List[List[str]]]:
         """Add user message to chat history and clear input."""
@@ -812,34 +813,34 @@ def create_application():
         # Clear the input
         return "", chat_history
     
-    async def bot_response(chat_history: List[List[str]]) -> List[List[str]]:
-        """Get bot response and update chat history."""
+    async def bot_response(chat_history: List[List[str]]) -> Tuple[List[List[str]], Optional[go.Figure]]:
+        """Get bot response and update chat history and return optional chart figure."""
         if not chat_history:
-            return chat_history
+            return chat_history, None
             
         # Get the last user message (first element of the last list if it exists)
         if not chat_history[-1][0] or chat_history[-1][1] is not None:
-            return chat_history
+            return chat_history, None
             
         try:
             question = chat_history[-1][0]
             logger.info(f"Processing question: {question}")
             
             # Call the agent and get the response
-            assistant_message = await stream_agent_response(question, chat_history[:-1])
+            assistant_message, chart_fig = await stream_agent_response(question, chat_history[:-1])
             
             # Update the assistant's message in the chat history
             chat_history[-1] = [question, assistant_message]
             
             logger.info("Response generation complete")
-            return chat_history
+            return chat_history, chart_fig
             
         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)
             if chat_history and len(chat_history[-1]) == 2 and chat_history[-1][1] is None:
                 chat_history[-1] = [chat_history[-1][0], error_msg]
-            return chat_history
+            return chat_history, None
     
     # Event handlers
     with demo:
@@ -852,7 +853,7 @@ def create_application():
         ).then(
             fn=bot_response,
             inputs=[chatbot],
-            outputs=[chatbot],
+            outputs=[chatbot, chart_display],
             api_name="ask"
         )
         
@@ -865,7 +866,7 @@ def create_application():
         ).then(
             fn=bot_response,
             inputs=[chatbot],
-            outputs=[chatbot]
+            outputs=[chatbot, chart_display]
         )
     
     return demo