import streamlit as st
import numpy as np
import pandas as pd
from smolagents import CodeAgent, tool
from typing import Union, List, Dict, Optional
import matplotlib.pyplot as plt
import seaborn as sns
import os
from groq import Groq
from dataclasses import dataclass
import tempfile
import base64
import io
import plotly.express as px
import plotly.graph_objects as go
# Set page configuration
st.set_page_config(
page_title="Data Analysis Assistant",
page_icon="📊",
layout="wide",
initial_sidebar_state="expanded"
)
# Custom CSS for DeepMind-inspired styling
st.markdown("""
""", unsafe_allow_html=True)
class GroqLLM:
"""Compatible LLM interface for smolagents CodeAgent"""
def __init__(self, model_name="llama-3.1-8B-Instant"):
self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
self.model_name = model_name
def __call__(self, prompt: Union[str, dict, List[Dict]]) -> str:
"""Make the class callable as required by smolagents"""
try:
# Handle different prompt formats
if isinstance(prompt, (dict, list)):
prompt_str = str(prompt)
else:
prompt_str = str(prompt)
# Create a properly formatted message
completion = self.client.chat.completions.create(
model=self.model_name,
messages=[{
"role": "user",
"content": prompt_str
}],
temperature=0.7,
max_tokens=1024,
stream=False
)
return completion.choices[0].message.content if completion.choices else "Error: No response generated"
except Exception as e:
error_msg = f"Error generating response: {str(e)}"
print(error_msg)
return error_msg
class DataAnalysisAgent(CodeAgent):
"""Extended CodeAgent with dataset awareness"""
def __init__(self, dataset: pd.DataFrame, *args, **kwargs):
super().__init__(*args, **kwargs)
self._dataset = dataset
@property
def dataset(self) -> pd.DataFrame:
"""Access the stored dataset"""
return self._dataset
def run(self, prompt: str) -> str:
"""Override run method to include dataset context"""
dataset_info = f"""
Dataset Shape: {self.dataset.shape}
Columns: {', '.join(self.dataset.columns)}
Data Types: {self.dataset.dtypes.to_dict()}
"""
enhanced_prompt = f"""
Analyze the following dataset:
{dataset_info}
Task: {prompt}
Use the provided tools to analyze this specific dataset and return detailed results.
"""
return super().run(enhanced_prompt)
@tool
def analyze_basic_stats(data: pd.DataFrame) -> str:
"""Calculate basic statistical measures for numerical columns in the dataset."""
# Access dataset from agent if no data provided
if data is None:
data = tool.agent.dataset
stats = {}
numeric_cols = data.select_dtypes(include=[np.number]).columns
for col in numeric_cols:
stats[col] = {
'mean': float(data[col].mean()),
'median': float(data[col].median()),
'std': float(data[col].std()),
'skew': float(data[col].skew()),
'missing': int(data[col].isnull().sum())
}
return str(stats)
@tool
def generate_correlation_matrix(data: pd.DataFrame) -> str:
"""Generate a visual correlation matrix for numerical columns in the dataset."""
# Access dataset from agent if no data provided
if data is None:
data = tool.agent.dataset
numeric_data = data.select_dtypes(include=[np.number])
# Using a modern Plotly heatmap instead of matplotlib
fig = px.imshow(
numeric_data.corr(),
text_auto=True,
aspect="auto",
color_continuous_scale="Blues",
title="Feature Correlation Matrix"
)
fig.update_layout(
height=600,
width=800,
font=dict(family="Inter, sans-serif"),
plot_bgcolor="white",
title_font=dict(size=20, color="#202124", family="Inter, sans-serif"),
margin=dict(l=40, r=40, t=60, b=40),
)
# Convert to HTML for display
fig_html = fig.to_html(full_html=False, include_plotlyjs='cdn')
return fig_html
@tool
def analyze_categorical_columns(data: pd.DataFrame) -> str:
"""Analyze categorical columns in the dataset for distribution and frequencies."""
# Access dataset from agent if no data provided
if data is None:
data = tool.agent.dataset
categorical_cols = data.select_dtypes(include=['object', 'category']).columns
analysis = {}
for col in categorical_cols:
analysis[col] = {
'unique_values': int(data[col].nunique()),
'top_categories': data[col].value_counts().head(5).to_dict(),
'missing': int(data[col].isnull().sum())
}
# Create an HTML visualization of categorical data
html_content = "
"
for col, stats in analysis.items():
html_content += f"
"
html_content += f"
{col}
"
html_content += f"
Unique Values: {stats['unique_values']}
"
html_content += f"
Missing Values: {stats['missing']}
"
# Add bar chart for top categories
if stats['top_categories']:
categories = list(stats['top_categories'].keys())
values = list(stats['top_categories'].values())
fig = go.Figure()
fig.add_trace(go.Bar(
x=categories,
y=values,
marker_color='#1a73e8',
hoverinfo='x+y'
))
fig.update_layout(
title=f"Top Categories for {col}",
xaxis_title="Category",
yaxis_title="Count",
font=dict(family="Inter, sans-serif"),
height=350,
margin=dict(l=40, r=40, t=60, b=80),
xaxis=dict(tickangle=-45)
)
html_content += fig.to_html(full_html=False, include_plotlyjs='cdn')
html_content += "
"
html_content += "
"
return html_content
@tool
def suggest_features(data: pd.DataFrame) -> str:
"""Suggest potential feature engineering steps based on data characteristics."""
# Access dataset from agent if no data provided
if data is None:
data = tool.agent.dataset
suggestions = []
numeric_cols = data.select_dtypes(include=[np.number]).columns
categorical_cols = data.select_dtypes(include=['object', 'category']).columns
if len(numeric_cols) >= 2:
suggestions.append("Consider creating interaction terms between numerical features")
if len(categorical_cols) > 0:
suggestions.append("Consider one-hot encoding for categorical variables")
for col in numeric_cols:
if data[col].skew() > 1 or data[col].skew() < -1:
suggestions.append(f"Consider log transformation for {col} due to skewness")
# Format as HTML for better display
html_content = """
Feature Engineering Suggestions
"""
for suggestion in suggestions:
html_content += f"""
✓{suggestion}
"""
if not suggestions:
html_content += """
!No specific feature engineering suggestions found for this dataset.
"""
html_content += """
"""
return html_content
@tool
def visualize_distributions(data: pd.DataFrame) -> str:
"""Create visualizations of numerical column distributions."""
# Access dataset from agent if no data provided
if data is None:
data = tool.agent.dataset
numeric_cols = data.select_dtypes(include=[np.number]).columns
if len(numeric_cols) == 0:
return "No numerical columns found in the dataset."
# Create HTML content with visualizations
html_content = "
"
# Create a grid of histograms using plotly
fig = make_subplots(rows=len(numeric_cols), cols=1,
subplot_titles=numeric_cols,
vertical_spacing=0.05)
for i, col in enumerate(numeric_cols):
fig.add_trace(
go.Histogram(
x=data[col].dropna(),
name=col,
marker_color='#1a73e8',
opacity=0.7
),
row=i+1, col=1
)
fig.update_layout(
height=300 * len(numeric_cols),
width=800,
title_text="Distribution of Numerical Features",
showlegend=False,
font=dict(family="Inter, sans-serif"),
margin=dict(l=40, r=40, t=40, b=20),
)
html_content += fig.to_html(full_html=False, include_plotlyjs='cdn')
html_content += "
"
return html_content
def generate_deepmind_logo():
"""Generate a placeholder logo similar to DeepMind's style."""
fig = go.Figure()
# Create simple geometric shapes for logo
fig.add_shape(
type="circle",
x0=0.3, y0=0.3, x1=0.7, y1=0.7,
line=dict(color="#1a73e8", width=3),
fillcolor="rgba(26, 115, 232, 0.2)",
)
fig.add_shape(
type="circle",
x0=0.4, y0=0.4, x1=0.6, y1=0.6,
line=dict(color="#1a73e8", width=2),
fillcolor="rgba(26, 115, 232, 0.4)",
)
fig.update_layout(
width=180,
height=60,
paper_bgcolor='rgba(0,0,0,0)',
plot_bgcolor='rgba(0,0,0,0)',
margin=dict(l=0, r=0, t=0, b=0),
showlegend=False,
xaxis=dict(showgrid=False, zeroline=False, visible=False),
yaxis=dict(showgrid=False, zeroline=False, visible=False),
)
return fig.to_html(full_html=False, include_plotlyjs='cdn')
def main():
# Logo and header
st.markdown("""
Data Analysis Assistant
Upload your dataset and get intelligent insights with AI-powered analysis
""", unsafe_allow_html=True)
# Initialize session state
if 'data' not in st.session_state:
st.session_state['data'] = None
if 'agent' not in st.session_state:
st.session_state['agent'] = None
if 'analysis_results' not in st.session_state:
st.session_state['analysis_results'] = None
# Create a two-column layout
col1, col2 = st.columns([1, 3])
with col1:
st.markdown('
', unsafe_allow_html=True)
st.markdown('
Upload Dataset
', unsafe_allow_html=True)
# File uploader with custom styling
uploaded_file = st.file_uploader("", type="csv")
if uploaded_file is not None:
try:
with st.spinner('Processing dataset...'):
# Load the dataset
data = pd.read_csv(uploaded_file)
st.session_state['data'] = data
# Initialize the agent with the dataset
st.session_state['agent'] = DataAnalysisAgent(
dataset=data,
tools=[analyze_basic_stats, generate_correlation_matrix,
analyze_categorical_columns, suggest_features,
visualize_distributions],
model=GroqLLM(),
additional_authorized_imports=["pandas", "numpy", "matplotlib",
"seaborn", "plotly"]
)
# Display dataset statistics
st.markdown("""
✓Dataset loaded successfully
""", unsafe_allow_html=True)
col1, col2 = st.columns(2)
with col1:
st.markdown(f"""
{data.shape[0]:,}
Rows
""", unsafe_allow_html=True)
with col2:
st.markdown(f"""
{data.shape[1]}
Columns
""", unsafe_allow_html=True)
except Exception as e:
st.error(f"Error: {str(e)}")
# Analysis type selection
if st.session_state['data'] is not None:
st.markdown('
', unsafe_allow_html=True)
# Main content area
with col2:
if st.session_state['data'] is not None:
# Data preview tab
st.markdown('
', unsafe_allow_html=True)
st.markdown('
Data Preview
', unsafe_allow_html=True)
# Add tabs for different data views
data_tabs = st.tabs(["Data Sample", "Column Info", "Missing Values"])
with data_tabs[0]:
st.markdown('
', unsafe_allow_html=True)
with data_tabs[1]:
col1, col2, col3 = st.columns(3)
with col1:
st.markdown("**Column Names**")
st.write(st.session_state['data'].columns.tolist())
with col2:
st.markdown("**Data Types**")
for col, dtype in st.session_state['data'].dtypes.items():
st.write(f"{col}: {dtype}")
with col3:
st.markdown("**Non-Null Count**")
for col, count in st.session_state['data'].count().items():
st.write(f"{col}: {count}/{len(st.session_state['data'])}")
with data_tabs[2]:
missing_data = st.session_state['data'].isnull().sum()
if missing_data.sum() > 0:
missing_df = pd.DataFrame({
'Column': missing_data.index,
'Missing Values': missing_data.values,
'Percentage': round(missing_data.values / len(st.session_state['data']) * 100, 2)
})
missing_df = missing_df[missing_df['Missing Values'] > 0].sort_values('Missing Values', ascending=False)
st.dataframe(missing_df, use_container_width=True)
# Add a visualization of missing values
fig = px.bar(
missing_df,
x='Column',
y='Percentage',
color='Percentage',
color_continuous_scale='Blues',
title='Missing Values by Column (%)'
)
fig.update_layout(
xaxis_title='',
yaxis_title='Missing Values (%)',
height=400
)
st.plotly_chart(fig, use_container_width=True)
else:
st.success("No missing values in the dataset!")
st.markdown('
', unsafe_allow_html=True)
# Analysis results section
if analysis_type:
st.markdown('
', unsafe_allow_html=True)
st.markdown(f'
{analysis_type} Results
', unsafe_allow_html=True)
if analysis_type == "Data Overview":
col1, col2 = st.columns(2)
with col1:
st.markdown("### Dataset Summary")
st.dataframe(st.session_state['data'].describe(), use_container_width=True)
with col2:
st.markdown("### Data Profile")
numeric_count = len(st.session_state['data'].select_dtypes(include=[np.number]).columns)
categorical_count = len(st.session_state['data'].select_dtypes(include=['object', 'category']).columns)
# Create a pie chart for data types
fig = px.pie(
values=[numeric_count, categorical_count],
names=['Numeric', 'Categorical'],
color_discrete_sequence=['#1a73e8', '#34a853'],
hole=0.4
)
fig.update_layout(
title='Column Types',
font=dict(family="Inter, sans-serif"),
legend=dict(orientation="h", yanchor="bottom", y=-0.2, xanchor="center", x=0.5)
)
st.plotly_chart(fig, use_container_width=True)
elif analysis_type == "Basic Statistics":
with st.spinner('Analyzing basic statistics...'):
result = st.session_state['agent'].run(
"Use the analyze_basic_stats tool to analyze this dataset and "
"provide insights about the numerical distributions."
)
# Parse the string representation of the dictionary
try:
# Remove the literal 'str' prefix if present
if result.startswith("str("):
result = result[4:-1]
# Convert string to dict
import ast
stats_dict = ast.literal_eval(result)
# Display results in a more visual format
for col, stats in stats_dict.items():
st.markdown(f"### {col}")
# Create metrics in columns
col1, col2, col3, col4 = st.columns(4)
with col1:
st.metric("Mean", f"{stats['mean']:.2f}")
with col2:
st.metric("Median", f"{stats['median']:.2f}")
with col3:
st.metric("Std Dev", f"{stats['std']:.2f}")
with col4:
st.metric("Skewness", f"{stats['skew']:.2f}")
# Create a boxplot for this column
fig = px.box(
st.session_state['data'],
y=col,
points="all",
color_discrete_sequence=['#1a73e8'],
title=f"Distribution of {col}"
)
fig.update_layout(
height=300,
margin=dict(t=40, b=20, l=40, r=20),
font=dict(family="Inter, sans-serif")
)
st.plotly_chart(fig, use_container_width=True)
st.markdown("---")
except Exception as e:
st.write(result)
elif analysis_type == "Feature Correlations":
with st.spinner('Analyzing feature correlations...'):
result = st.session_state['agent'].run(
"Use the generate_correlation_matrix tool to analyze correlations "
"and explain any strong relationships found."
)
# If the result is HTML, display it directly
if isinstance(result, str) and ("', unsafe_allow_html=True)
else:
# Display welcome message for users who haven't uploaded data yet
st.markdown("""
Welcome to Data Analysis Assistant
Upload a CSV file to get started with instant insights and intelligent analysis.
Our AI-powered assistant will help you understand your data like never before.
📊
Automatic Visualizations
Get instant charts and plots revealing insights in your data
ðŸ§
AI-Powered Analysis
Advanced algorithms find patterns and correlations automatically
💡
Smart Recommendations
Get suggestions for feature engineering and data preparation
""", unsafe_allow_html=True)
# Import for subplot creation
from plotly.subplots import make_subplots
if __name__ == "__main__":
# Check if Groq API key is available
if not os.environ.get("GROQ_API_KEY"):
st.error("""
GROQ API key not found! Please set your GROQ_API_KEY environment variable.
You can get an API key from https://console.groq.com/
""")
else:
main()