Update app.py

app.py

@@ -12,76 +12,60 @@ from typing import Dict, List, Optional
 from langchain.tools import tool
 from langchain.agents import initialize_agent, AgentType
 from scipy.stats import ttest_ind, f_oneway
-from statsmodels.tsa.seasonal import seasonal_decompose
-from statsmodels.tsa.stattools import adfuller
-from jinja2 import Template
 
 # Initialize Groq Client
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 
-
 class ResearchInput(BaseModel):
-    """Base schema for research tool inputs, ensuring type and description integrity."""
-    data_key: str = Field(..., description="Session state key containing the DataFrame.")
-    columns: Optional[List[str]] = Field(None, description="List of column names to analyze.")
-
+    """Base schema for research tool inputs"""
+    data_key: str = Field(..., description="Session state key containing DataFrame")
+    columns: Optional[List[str]] = Field(None, description="List of columns to analyze")
 
 class TemporalAnalysisInput(ResearchInput):
-    """Schema for temporal analysis inputs, focusing on specific time-series requirements."""
-    time_col: str = Field(..., description="Name of the column containing timestamp data.")
-    value_col: str = Field(..., description="Name of the column containing numerical values to analyze.")
-
+    """Schema for temporal analysis"""
+    time_col: str = Field(..., description="Name of timestamp column")
+    value_col: str = Field(..., description="Name of value column to analyze")
 
 class HypothesisInput(ResearchInput):
-    """Schema for hypothesis testing, demanding group and value specification for statistical rigor."""
-    group_col: str = Field(..., description="Categorical column defining the groups for comparison.")
-    value_col: str = Field(..., description="Numerical column for comparing means across groups.")
-
+    """Schema for hypothesis testing"""
+    group_col: str = Field(..., description="Categorical column defining groups")
+    value_col: str = Field(..., description="Numerical column to compare")
 
 class GroqResearcher:
-    """
-    A sophisticated AI research engine powered by Groq, designed for rigorous academic-style analysis.
-    This class handles complex data queries and delivers structured research outputs.
-    """
-
+    """Advanced AI Research Engine using Groq"""
     def __init__(self, model_name="mixtral-8x7b-32768"):
         self.model_name = model_name
-        self.system_template = """
-        You are a senior data scientist at a prestigious research institution. Your analysis must 
-        adhere to rigorous scientific standards. Consider the dataset properties and the user's query.
+        self.system_template = """You are a senior data scientist at a research institution. 
+        Analyze this dataset with rigorous statistical methods and provide academic-quality insights:
+        {dataset_info}
         
-        Dataset Context:
-        - Dimensions: {{ dataset_shape }}
-        - Variables: {{ dataset_variables }}
-        - Temporal Coverage: {{ temporal_coverage }}
-        - Missing Value Counts: {{ missing_values }}
-
-        User Inquiry: {{ query }}
-
-        Response Structure (Critical for all analyses):
-        1. **Executive Summary:** Provide a 1-2 paragraph overview of the findings, contextualized within the dataset's characteristics.
-        2. **Methodology:** Detail the exact analysis techniques used, including statistical tests or model types, and their justification.
-        3. **Key Findings:** Present the most significant observations and statistical results (p-values, effect sizes) with proper interpretation.
-        4. **Limitations:** Acknowledge and describe the constraints of the dataset or analytical methods that might affect the results' interpretation or generalizability.
-        5. **Recommended Next Steps:** Suggest future studies, experiments, or analyses that could extend the current investigation and address the noted limitations.
-
-        """
+        User Question: {query}
+        
+        Required Format:
+        - Executive Summary (1 paragraph)
+        - Methodology (bullet points)
+        - Key Findings (numbered list)
+        - Limitations
+        - Recommended Next Steps"""
 
     def research(self, query: str, data: pd.DataFrame) -> str:
-        """Executes in-depth research using the Groq API to produce academic-quality analyses."""
+        """Conduct academic-level analysis using Groq"""
         try:
-            dataset_info = {
-                "dataset_shape": str(data.shape),
-                "dataset_variables": ", ".join(data.columns),
-                "temporal_coverage": str(data.select_dtypes(include='datetime').columns.tolist()),
-                "missing_values": str(data.isnull().sum().to_dict()),
-            }
-
-            prompt = Template(self.system_template).render(**dataset_info, query=query)
-
+            dataset_info = f"""
+            Dataset Dimensions: {data.shape}
+            Variables: {', '.join(data.columns)}
+            Temporal Coverage: {data.select_dtypes(include='datetime').columns.tolist()}
+            Missing Values: {data.isnull().sum().to_dict()}
+            """
+            
+            prompt = PromptTemplate.from_template(self.system_template).format(
+                dataset_info=dataset_info,
+                query=query
+            )
+            
             completion = client.chat.completions.create(
                 messages=[
-                    {"role": "system", "content": "You are a research AI assistant."},
+                    {"role": "system", "content": "You are a research AI assistant"},
                     {"role": "user", "content": prompt}
                 ],
                 model=self.model_name,
@@ -89,22 +73,20 @@ class GroqResearcher:
                 max_tokens=4096,
                 stream=False
             )
+            
             return completion.choices[0].message.content
+        
         except Exception as e:
-            return f"Research Error Encountered: {str(e)}"
-
+            return f"Research Error: {str(e)}"
 
 @tool(args_schema=ResearchInput)
 def advanced_eda(data_key: str) -> Dict:
-    """
-    Performs a comprehensive Exploratory Data Analysis, including statistical profiling,
-    temporal analysis of datetime columns, and detailed quality checks.
-    """
+    """Comprehensive Exploratory Data Analysis with Statistical Profiling"""
     try:
         data = st.session_state[data_key]
         analysis = {
             "dimensionality": {
-                "rows": int(len(data)),  # Ensure rows are an integer
+                "rows": len(data),
                 "columns": list(data.columns),
                 "memory_usage": f"{data.memory_usage().sum() / 1e6:.2f} MB"
             },
@@ -112,147 +94,112 @@ def advanced_eda(data_key: str) -> Dict:
             "temporal_analysis": {
                 "date_ranges": {
                     col: {
-                        "min": str(data[col].min()),  # Ensure date is a string
-                        "max": str(data[col].max())  # Ensure date is a string
+                        "min": data[col].min(),
+                        "max": data[col].max()
                     } for col in data.select_dtypes(include='datetime').columns
                 }
             },
             "data_quality": {
                 "missing_values": data.isnull().sum().to_dict(),
-                "duplicates": int(data.duplicated().sum()),  # Ensure duplicates are an integer
+                "duplicates": data.duplicated().sum(),
                 "cardinality": {
-                    col: int(data[col].nunique()) for col in data.columns  # Ensure cardinality is integer
+                    col: data[col].nunique() for col in data.columns
                 }
             }
         }
         return analysis
     except Exception as e:
-        return {"error": f"Advanced EDA Failed: {str(e)}"}
+        return {"error": f"EDA Failed: {str(e)}"}
 
 @tool(args_schema=ResearchInput)
 def visualize_distributions(data_key: str, columns: List[str]) -> str:
-    """
-    Generates high-quality, publication-ready distribution visualizations (histograms with KDE)
-    for selected numerical columns, and returns the image as a base64 encoded string.
-    """
+    """Generate publication-quality distribution visualizations"""
     try:
         data = st.session_state[data_key]
-        plt.figure(figsize=(15, 7))  # Adjusted figure size for better readability
+        plt.figure(figsize=(12, 6))
         for i, col in enumerate(columns, 1):
             plt.subplot(1, len(columns), i)
-            sns.histplot(data[col], kde=True, stat="density", color=sns.color_palette()[i % len(sns.color_palette())])
-            plt.title(f'Distribution of {col}', fontsize=14, fontweight='bold')  # Enhanced title
-            plt.xlabel(col, fontsize=12)
-            plt.ylabel('Density', fontsize=12)
-            plt.xticks(fontsize=10)
-            plt.yticks(fontsize=10)
-            plt.grid(axis='y', linestyle='--')
-            sns.despine(top=True, right=True)  # Improved styling
-        plt.tight_layout(pad=2)  # Added padding for tight layout
+            sns.histplot(data[col], kde=True, stat="density")
+            plt.title(f'Distribution of {col}', fontsize=10)
+            plt.xticks(fontsize=8)
+            plt.yticks(fontsize=8)
+        plt.tight_layout()
         
         buf = io.BytesIO()
         plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
         plt.close()
         return base64.b64encode(buf.getvalue()).decode()
     except Exception as e:
-        return f"Distribution Visualization Error: {str(e)}"
-
+        return f"Visualization Error: {str(e)}"
 
 @tool(args_schema=TemporalAnalysisInput)
 def temporal_analysis(data_key: str, time_col: str, value_col: str) -> Dict:
-    """
-    Performs a sophisticated time series analysis, including decomposition and trend assessment,
-    providing both statistical insights and a visual representation.
-    """
+    """Time Series Decomposition and Trend Analysis"""
     try:
         data = st.session_state[data_key]
-        ts_data = data.set_index(pd.to_datetime(data[time_col]))[value_col].dropna()  # Handle NaNs
-
-        if ts_data.empty:
-            return {"error": "No valid time series data found for analysis after NaN removal."}
-
-        decomposition = seasonal_decompose(ts_data, model='additive', period=min(len(ts_data), 365) if len(ts_data) > 10 else 1)
-         
-        plt.figure(figsize=(16, 10))
+        ts_data = data.set_index(pd.to_datetime(data[time_col]))[value_col]
+        
+        decomposition = seasonal_decompose(ts_data, period=365)
+        
+        plt.figure(figsize=(12, 8))
         decomposition.plot()
         plt.tight_layout()
-
+        
         buf = io.BytesIO()
-        plt.savefig(buf, format='png', dpi=300)  # Increased dpi for higher resolution
+        plt.savefig(buf, format='png')
         plt.close()
         plot_data = base64.b64encode(buf.getvalue()).decode()
-         
-        adf_result = adfuller(ts_data)
-        stationarity_p_value = adf_result[1]
-
+        
         return {
             "trend_statistics": {
-                "stationarity": stationarity_p_value,
-                "stationarity_interpretation": interpret_p_value(stationarity_p_value),
-                "seasonality_strength": max(decomposition.seasonal) if hasattr(decomposition, 'seasonal') else None
+                "stationarity": adfuller(ts_data)[1],
+                "seasonality_strength": max(decomposition.seasonal)
             },
-            "visualization": plot_data,
-            "decomposition_data": {
-                "trend": decomposition.trend.dropna().to_dict() if hasattr(decomposition, 'trend') else None,
-                "seasonal": decomposition.seasonal.dropna().to_dict() if hasattr(decomposition, 'seasonal') else None,
-                "residual": decomposition.resid.dropna().to_dict() if hasattr(decomposition, 'resid') else None,
-            }
+            "visualization": plot_data
         }
     except Exception as e:
-        return {"error": f"Temporal Analysis Failure: {str(e)}"}
+        return {"error": f"Temporal Analysis Failed: {str(e)}"}
 
 @tool(args_schema=HypothesisInput)
 def hypothesis_testing(data_key: str, group_col: str, value_col: str) -> Dict:
-    """
-    Conducts statistical hypothesis testing, providing detailed test results, effect size measures,
-    and interpretations for both t-tests and ANOVAs.
-    """
+    """Statistical Hypothesis Testing with Automated Assumption Checking"""
     try:
         data = st.session_state[data_key]
         groups = data[group_col].unique()
         
         if len(groups) < 2:
-            return {"error": "Insufficient groups for comparison. Must have at least two groups."}
-            
-        group_data = [data[data[group_col] == g][value_col].dropna() for g in groups]
-
-        if any(len(group) < 2 for group in group_data):
-             return {"error": "Each group must have at least two data points for testing."}
+            return {"error": "Insufficient groups for comparison"}
             
         if len(groups) == 2:
+            group_data = [data[data[group_col] == g][value_col] for g in groups]
             stat, p = ttest_ind(*group_data)
             test_type = "Independent t-test"
         else:
+            group_data = [data[data[group_col] == g][value_col] for g in groups]
             stat, p = f_oneway(*group_data)
             test_type = "ANOVA"
         
-        effect_size = None
-        if len(groups) == 2:
-            pooled_variance = np.sqrt((group_data[0].var() + group_data[1].var()) / 2)
-            if pooled_variance != 0:
-                cohens_d = abs(group_data[0].mean() - group_data[1].mean()) / pooled_variance
-                effect_size = {"cohens_d": cohens_d}
-            else:
-                effect_size = {"cohens_d": None, "error": "Cannot compute effect size due to zero pooled variance."}
-
         return {
             "test_type": test_type,
-            "test_statistic": float(stat),  # Ensure stat is a float
-            "p_value": float(p),  # Ensure p_value is a float
-            "effect_size": effect_size,
-            "interpretation": interpret_p_value(p),
-             "group_means": {g: float(data[data[group_col] == g][value_col].mean()) for g in groups} # Group Means
+            "test_statistic": stat,
+            "p_value": p,
+            "effect_size": {
+                "cohens_d": abs(group_data[0].mean() - group_data[1].mean())/np.sqrt(
+                    (group_data[0].var() + group_data[1].var())/2
+                ) if len(groups) == 2 else None
+            },
+            "interpretation": interpret_p_value(p)
         }
     except Exception as e:
         return {"error": f"Hypothesis Testing Failed: {str(e)}"}
 
 def interpret_p_value(p: float) -> str:
-    """Provides nuanced interpretations of p-values, including qualitative descriptors."""
-    if p < 0.001: return "Highly significant evidence against the null hypothesis (p < 0.001)."
-    elif p < 0.01: return "Strong evidence against the null hypothesis (0.001 ≤ p < 0.01)."
-    elif p < 0.05: return "Moderate evidence against the null hypothesis (0.01 ≤ p < 0.05)."
-    elif p < 0.1: return "Weak evidence against the null hypothesis (0.05 ≤ p < 0.1)."
-    else: return "No significant evidence against the null hypothesis (p ≥ 0.1)."
+    """Scientific interpretation of p-values"""
+    if p < 0.001: return "Very strong evidence against H0"
+    elif p < 0.01: return "Strong evidence against H0"
+    elif p < 0.05: return "Evidence against H0"
+    elif p < 0.1: return "Weak evidence against H0"
+    else: return "No significant evidence against H0"
 
 def main():
     st.set_page_config(page_title="AI Research Lab", layout="wide")
@@ -270,12 +217,9 @@ def main():
         uploaded_file = st.file_uploader("Upload research dataset", type=["csv", "parquet"])
         if uploaded_file:
             with st.spinner("Initializing dataset..."):
-                try:
-                    st.session_state.data = pd.read_csv(uploaded_file)
-                    st.success(f"Loaded {len(st.session_state.data):,} research observations")
-                except Exception as e:
-                   st.error(f"Error loading the dataset. Please ensure it's a valid CSV or Parquet format. Error details: {e}")
-        
+                st.session_state.data = pd.read_csv(uploaded_file)
+                st.success(f"Loaded {len(st.session_state.data):,} research observations")
+    
     # Main research interface
     if st.session_state.data is not None:
         col1, col2 = st.columns([1, 3])
@@ -286,10 +230,10 @@ def main():
                 "Variables": list(st.session_state.data.columns),
                 "Time Range": {
                     col: {
-                        "min": str(st.session_state.data[col].min()),
-                        "max": str(st.session_state.data[col].max())
+                        "min": st.session_state.data[col].min(),
+                        "max": st.session_state.data[col].max()
                     } for col in st.session_state.data.select_dtypes(include='datetime').columns
-                } if st.session_state.data.select_dtypes(include='datetime').columns.tolist() else "No Temporal Data",
+                },
                 "Size": f"{st.session_state.data.memory_usage().sum() / 1e6:.2f} MB"
             })
         
@@ -310,42 +254,35 @@ def main():
                     st.json(eda_result)
                 
                 elif analysis_type == "Temporal Pattern Analysis":
-                    time_cols = st.session_state.data.select_dtypes(include='datetime').columns.tolist()
-                    if not time_cols:
-                        st.warning("No datetime columns detected. Please ensure you have a datetime column for this analysis.")
-                    else:
-                        time_col = st.selectbox("Temporal Variable", time_cols)
-                        value_col = st.selectbox("Analysis Variable",
-                            st.session_state.data.select_dtypes(include=np.number).columns)
-                        
-                        if time_col and value_col:
-                            result = temporal_analysis.invoke({
-                                "data_key": "data",
-                                "time_col": time_col,
-                                "value_col": value_col
-                            })
-                            if "visualization" in result:
-                                st.image(f"data:image/png;base64,{result['visualization']}",
-                                    use_column_width=True)
-                            st.json(result)
+                    time_col = st.selectbox("Temporal Variable", 
+                        st.session_state.data.select_dtypes(include='datetime').columns)
+                    value_col = st.selectbox("Analysis Variable",
+                        st.session_state.data.select_dtypes(include=np.number).columns)
+                    
+                    if time_col and value_col:
+                        result = temporal_analysis.invoke({
+                            "data_key": "data",
+                            "time_col": time_col,
+                            "value_col": value_col
+                        })
+                        if "visualization" in result:
+                            st.image(f"data:image/png;base64,{result['visualization']}")
+                        st.json(result)
                 
                 elif analysis_type == "Comparative Statistics":
-                    cat_cols = st.session_state.data.select_dtypes(include='category').columns.tolist() + st.session_state.data.select_dtypes(include='object').columns.tolist()
-                    if not cat_cols:
-                         st.warning("No categorical columns detected. Please ensure you have a categorical column for this analysis.")
-                    else:
-                        group_col = st.selectbox("Grouping Variable", cat_cols)
-                        value_col = st.selectbox("Metric Variable",
-                            st.session_state.data.select_dtypes(include=np.number).columns)
-                        
-                        if group_col and value_col:
-                            result = hypothesis_testing.invoke({
-                                "data_key": "data",
-                                "group_col": group_col,
-                                "value_col": value_col
-                            })
-                            st.subheader("Statistical Test Results")
-                            st.json(result)
+                    group_col = st.selectbox("Grouping Variable",
+                        st.session_state.data.select_dtypes(include='category').columns)
+                    value_col = st.selectbox("Metric Variable",
+                        st.session_state.data.select_dtypes(include=np.number).columns)
+                    
+                    if group_col and value_col:
+                        result = hypothesis_testing.invoke({
+                            "data_key": "data",
+                            "group_col": group_col,
+                            "value_col": value_col
+                        })
+                        st.subheader("Statistical Test Results")
+                        st.json(result)
                 
                 elif analysis_type == "Distribution Analysis":
                     num_cols = st.session_state.data.select_dtypes(include=np.number).columns.tolist()
@@ -355,8 +292,7 @@ def main():
                             "data_key": "data",
                             "columns": selected_cols
                         })
-                        st.image(f"data:image/png;base64,{img_data}",
-                                 use_column_width=True)
+                        st.image(f"data:image/png;base64,{img_data}")
             
             with research_tab:
                 research_query = st.text_area("Enter Research Question:", height=150,