Error Fix

app.py

@@ -300,16 +300,9 @@ def process_text_query(query, history):
                     try:
                         print("Visualization requested, attempting to create plot...")
                         
-                        # Improved visualization type detection
-                        viz_keywords = {
-                            'bar': ['bar chart', 'bar graph', 'bar plot', 'barchart', 'bargraph'],
-                            'line': ['line chart', 'line graph', 'line plot', 'linechart', 'trend', 'trends', 'time series'],
-                            'pie': ['pie chart', 'pie graph', 'pie plot', 'piechart', 'distribution', 'proportion'],
-                            'histogram': ['histogram', 'distribution of', 'frequency distribution'],
-                            'box': ['box plot', 'boxplot', 'box and whisker', 'outliers', 'quartiles'],
-                            'heatmap': ['heatmap', 'heat map', 'correlation matrix', 'correlation heatmap'],
-                            'scatter': ['scatter', 'scatter plot', 'relationship between', 'correlation between']
-                        }
+                        # Increase plot size
+                        fig_width = 1000  # Increased from 900
+                        fig_height = 700  # Increased from 600
                         
                         # Determine visualization type from query
                         viz_type = None
@@ -334,10 +327,6 @@ def process_text_query(query, history):
                         
                         # Create the appropriate visualization based on type
                         if len(numeric_cols) >= 1 and len(result_df) > 1:
-                            # Set common figure parameters
-                            fig_width = 900
-                            fig_height = 600
-                            
                             if viz_type == 'pie' and len(result_df) <= 20:
                                 # For pie charts, we need a category column and a value column
                                 category_col = result_df.columns[0]
@@ -357,35 +346,72 @@ def process_text_query(query, history):
                                     color_discrete_sequence=px.colors.qualitative.Pastel
                                 )
                                 
-                            elif viz_type == 'histogram' and numeric_cols:
-                                # For histograms, manually create the bins if not already binned
+                            elif viz_type == 'histogram' and len(result_df.columns) > 0:
+                                # For histograms, we need at least one column
+                                
+                                # Find the best column for histogram (prefer numeric)
+                                if numeric_cols:
                                 x_col = numeric_cols[0]
+                                else:
+                                    x_col = result_df.columns[0]
                                 
-                                if len(result_df[x_col].unique()) > 20:
-                                    # Data is not pre-binned, create a histogram
+                                # Check if data is already binned
+                                if len(result_df) <= 30 and 'bin' in result_df.columns or 'range' in result_df.columns:
+                                    # Data is pre-binned, use a bar chart
+                                    bin_col = 'bin' if 'bin' in result_df.columns else 'range'
+                                    count_col = 'count' if 'count' in result_df.columns else numeric_cols[0] if numeric_cols else result_df.columns[1]
+                                    
+                                    fig = px.bar(
+                                        result_df,
+                                        x=bin_col,
+                                        y=count_col,
+                                        title=f"Histogram of {x_col}",
+                                        labels={bin_col: x_col, count_col: 'Frequency'},
+                                        color_discrete_sequence=['#636EFA']
+                                    )
+                                else:
+                                    # Create a proper histogram from raw data
                                     fig = px.histogram(
                                         result_df, 
                                         x=x_col,
                                         title=f"Distribution of {x_col}",
                                         nbins=20,
                                         marginal="box",  # Add a box plot on the margin
-                                        color_discrete_sequence=['#636EFA']
+                                        color_discrete_sequence=['#636EFA'],
+                                        opacity=0.8,
+                                        histnorm='probability density'  # Normalize to show density instead of count
                                     )
-                                else:
-                                    # Data is likely pre-binned, use a bar chart
-                                    x_col = result_df.columns[0]
-                                    y_col = numeric_cols[0] if x_col not in numeric_cols else numeric_cols[1] if len(numeric_cols) > 1 else 'count'
                                     
-                                    if y_col == 'count' and 'count' not in result_df.columns:
-                                        # If we need a count column but don't have one
-                                        result_df['count'] = 1
+                                    # Add a KDE (kernel density estimate) curve
+                                    from scipy import stats
+                                    import numpy as np
                                     
-                                    fig = px.bar(
-                                        result_df,
-                                        x=x_col,
-                                        y=y_col,
-                                        title=f"Histogram of {x_col}",
-                                        color_discrete_sequence=['#636EFA']
+                                    # Only add KDE if we have numeric data
+                                    if pd.api.types.is_numeric_dtype(result_df[x_col]):
+                                        # Remove NaN values
+                                        data = result_df[x_col].dropna()
+                                        
+                                        if len(data) > 1:  # Need at least 2 points for KDE
+                                            # Calculate KDE
+                                            kde = stats.gaussian_kde(data)
+                                            x_range = np.linspace(data.min(), data.max(), 1000)
+                                            y_kde = kde(x_range)
+                                            
+                                            # Add KDE curve
+                                            fig.add_scatter(
+                                                x=x_range, 
+                                                y=y_kde,
+                                                mode='lines',
+                                                line=dict(color='red', width=2),
+                                                name='Density Curve'
+                                            )
+                                
+                                # Improve histogram layout
+                                fig.update_layout(
+                                    bargap=0.1,  # Gap between bars
+                                    xaxis_title=x_col,
+                                    yaxis_title='Frequency',
+                                    showlegend=True
                                     )
                             
                             elif viz_type == 'box' and numeric_cols:
@@ -416,7 +442,14 @@ def process_text_query(query, history):
                                 # If we have many numeric columns, create a correlation matrix
                                 if len(numeric_cols) >= 3:
                                     # Create a correlation matrix
-                                    corr_df = result_df[numeric_cols].corr()
+                                    # First, drop any rows with NaN values in numeric columns
+                                    clean_df = result_df[numeric_cols].dropna()
+                                    
+                                    if len(clean_df) > 1:  # Need at least 2 rows for correlation
+                                        corr_df = clean_df.corr()
+                                        
+                                        # Round to 2 decimal places for display
+                                        corr_df = corr_df.round(2)
                                     
                                     fig = px.imshow(
                                         corr_df,
@@ -426,13 +459,31 @@ def process_text_query(query, history):
                                         aspect="auto",
                                         zmin=-1, zmax=1  # Set limits for correlation values
                                     )
+                                        
+                                        # Improve heatmap layout
+                                fig.update_layout(
+                                            xaxis_title="Features",
+                                            yaxis_title="Features",
+                                            coloraxis_colorbar=dict(
+                                                title="Correlation",
+                                                thicknessmode="pixels", thickness=20,
+                                                lenmode="pixels", len=300,
+                                                yanchor="top", y=1,
+                                                ticks="outside"
+                                            )
+                                        )
+                                    else:
+                                        # Not enough data for correlation
+                                        fig = px.bar(
+                                            pd.DataFrame({'Message': ['Not enough data for heatmap']}),
+                                            title="Cannot create heatmap - insufficient data"
+                                        )
                                 else:
-                                    # If we only have 2 numeric columns, we need to bin the data
-                                    # to create a 2D histogram (heatmap)
+                                    # If we only have 2 numeric columns, create a 2D histogram
                                     x_col = numeric_cols[0]
                                     y_col = numeric_cols[1]
                                     
-                                    # Create a 2D histogram
+                                    # Create a 2D histogram (heatmap)
                                     fig = px.density_heatmap(
                                         result_df,
                                         x=x_col,
@@ -440,7 +491,22 @@ def process_text_query(query, history):
                                         title=f"Density Heatmap of {x_col} vs {y_col}",
                                         color_continuous_scale='Viridis',
                                         nbinsx=20,
-                                        nbinsy=20
+                                        nbinsy=20,
+                                        marginal_x="histogram",  # Add histograms on the margins
+                                        marginal_y="histogram"
+                                    )
+                                    
+                                    # Improve heatmap layout
+                                    fig.update_layout(
+                                        xaxis_title=x_col,
+                                        yaxis_title=y_col,
+                                        coloraxis_colorbar=dict(
+                                            title="Count",
+                                            thicknessmode="pixels", thickness=20,
+                                            lenmode="pixels", len=300,
+                                            yanchor="top", y=1,
+                                            ticks="outside"
+                                        )
                                     )
                             
                             elif viz_type == 'scatter' and len(numeric_cols) >= 2:
@@ -451,29 +517,56 @@ def process_text_query(query, history):
                                 # Add a third dimension (size) if available
                                 size_col = numeric_cols[2] if len(numeric_cols) > 2 else None
                                 
+                                # Add a color dimension if available
+                                if len(result_df.columns) > len(numeric_cols):
+                                    # Find a categorical column for color
+                                    categorical_cols = [col for col in result_df.columns if col not in numeric_cols]
+                                    color_col = categorical_cols[0] if categorical_cols else None
+                                else:
+                                    color_col = None
+                                
+                                # Create scatter plot with enhanced features
                                 fig = px.scatter(
                                     result_df,
                                     x=x_col,
                                     y=y_col,
                                     size=size_col,
+                                    color=color_col,  # Add color dimension if available
                                     title=f"Relationship between {x_col} and {y_col}",
                                     opacity=0.7,
-                                    color_discrete_sequence=['#636EFA']
+                                    size_max=15,  # Maximum marker size
+                                    color_discrete_sequence=px.colors.qualitative.Plotly
                                 )
                                 
                                 # Add a trend line
+                                if pd.api.types.is_numeric_dtype(result_df[x_col]) and pd.api.types.is_numeric_dtype(result_df[y_col]):
+                                    fig.update_layout(
+                                        shapes=[
+                                            dict(
+                                                type='line',
+                                                xref='x', yref='y',
+                                                x0=result_df[x_col].min(),
+                                                y0=result_df[y_col].min(),
+                                                x1=result_df[x_col].max(),
+                                                y1=result_df[y_col].max(),
+                                                line=dict(color='red', width=2, dash='dash')
+                                            )
+                                        ]
+                                    )
+                                
+                                # Improve scatter plot layout
                                 fig.update_layout(
-                                    shapes=[
-                                        dict(
-                                            type='line',
-                                            xref='x', yref='y',
-                                            x0=result_df[x_col].min(),
-                                            y0=result_df[y_col].min(),
-                                            x1=result_df[x_col].max(),
-                                            y1=result_df[y_col].max(),
-                                            line=dict(color='red', width=2, dash='dash')
-                                        )
-                                    ]
+                                    xaxis_title=x_col,
+                                    yaxis_title=y_col,
+                                    showlegend=True,
+                                    legend=dict(
+                                        title=color_col if color_col else "",
+                                        orientation="h",
+                                        yanchor="bottom",
+                                        y=1.02,
+                                        xanchor="right",
+                                        x=1
+                                    )
                                 )
                             
                             elif viz_type == 'line':
@@ -539,7 +632,9 @@ def process_text_query(query, history):
                                 margin=dict(l=50, r=50, b=100, t=100, pad=4),
                                 template="plotly_white",
                                 font=dict(size=14),
-                                legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1)
+                                legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
+                                plot_bgcolor='rgba(240,240,240,0.2)',  # Light gray background
+                                paper_bgcolor='white'
                             )
                             
                             # Convert the figure to an image and encode it as base64
@@ -547,8 +642,8 @@ def process_text_query(query, history):
                             encoded = base64.b64encode(img_bytes).decode("ascii")
                             img_src = f"data:image/png;base64,{encoded}"
                             
-                            # Add the image directly to the response
-                            response += f"\n\n<img src='{img_src}' width='100%' />"
+                            # Add the image directly to the response with increased size
+                            response += f"\n\n<img src='{img_src}' width='100%' style='min-height:700px;' />"
                             
                             # Add note about visualization
                             response += f"\n\n**A {viz_type} visualization has been generated and is displayed above.**"
@@ -880,7 +975,7 @@ with gr.Blocks(title="AI Document Analysis & Voice Assistant") as demo:
                     show_label=False
                 )
             with gr.Column(scale=1):
-                voice_btn = gr.Button("🎤")
+                voice_btn = gr.Button("��")
         
         with gr.Row():
             submit_btn = gr.Button("Submit")