v1

app.py

@@ -3,8 +3,15 @@ import sys
 import re
 import gradio as gr
 import json
-from typing import List, Dict, Any, Optional, Tuple
+import tempfile
+import base64
+import io
+from typing import List, Dict, Any, Optional, Tuple, Union
 import logging
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
 try:
     # Intentar importar dependencias opcionales
@@ -29,6 +36,70 @@ logger = logging.getLogger(__name__)
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
+
+def generate_chart(data: Union[Dict, List[Dict], pd.DataFrame], 
+                 chart_type: str, 
+                 x: str, 
+                 y: str = None, 
+                 title: str = "", 
+                 x_label: str = None, 
+                 y_label: str = None) -> str:
+    """
+    Generate a chart from data and return it as an HTML string with embedded image.
+    
+    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)
+        title: Chart title
+        x_label: Label for x-axis
+        y_label: Label for y-axis
+        
+    Returns:
+        HTML string with embedded chart image
+    """
+    try:
+        # Convert data to DataFrame if it's a list of dicts
+        if isinstance(data, list):
+            df = pd.DataFrame(data)
+        elif isinstance(data, dict):
+            df = pd.DataFrame([data])
+        else:
+            df = data
+            
+        if not isinstance(df, pd.DataFrame):
+            return "Error: Data must be a dictionary, list of dictionaries, or pandas DataFrame"
+            
+        # Generate the appropriate chart type
+        if chart_type == 'bar':
+            fig = px.bar(df, x=x, y=y, title=title)
+        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)
+        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'"
+        
+        # Update layout
+        fig.update_layout(
+            xaxis_title=x_label or x,
+            yaxis_title=y_label or (y if y != x else ''),
+            title=title or f"{chart_type.capitalize()} Chart of {x} vs {y}" if y else f"{chart_type.capitalize()} Chart of {x}",
+            template="plotly_white"
+        )
+        
+        # Convert to HTML
+        return f"<div style='width: 100%;'>{fig.to_html(full_html=False)}</div>"
+        
+    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>"
 logger = logging.getLogger(__name__)
 
 def check_environment():
@@ -400,7 +471,61 @@ async def stream_agent_response(question: str, chat_history: List) -> List[Dict]
                 db_connection, _ = setup_database_connection()
                 if db_connection:
                     query_result = execute_sql_query(sql_query, db_connection)
+                    
+                    # 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:
+                        if isinstance(query_result, str) and '|' in query_result and '---' in query_result:
+                            # Convert markdown table to DataFrame
+                            from io import StringIO
+                            import re
+                            
+                            # Clean up the markdown table
+                            lines = [line.strip() for line in query_result.split('\n') 
+                                    if line.strip() and '---' not in line and '|' in line]
+                            if len(lines) > 1:  # At least header + 1 data row
+                                # Get column names from the first line
+                                columns = [col.strip() for col in lines[0].split('|')[1:-1]]
+                                # Get data rows
+                                data = []
+                                for line in lines[1:]:
+                                    values = [val.strip() for val in line.split('|')[1:-1]]
+                                    if len(values) == len(columns):
+                                        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}"
+                    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
+                        response_text += "\n\n⚠️ No se pudo generar la visualización de los datos."
                 else:
                     response_text += "\n\n⚠️ No se pudo conectar a la base de datos para ejecutar la consulta."
             

requirements.txt

@@ -12,3 +12,5 @@ python-dotenv==1.0.1
 pymysql==1.1.0
 numpy==1.26.4
 python-multipart>=0.0.18  # Required by gradio
+plotly==5.18.0  # For interactive charts
+kaleido==0.2.1  # For saving plotly charts as images