Update app.py

app.py

@@ -12,60 +12,76 @@ 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"""
-    data_key: str = Field(..., description="Session state key containing DataFrame")
-    columns: Optional[List[str]] = Field(None, description="List of columns to analyze")
+    """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.")
+
 
 class TemporalAnalysisInput(ResearchInput):
-    """Schema for temporal analysis"""
-    time_col: str = Field(..., description="Name of timestamp column")
-    value_col: str = Field(..., description="Name of value column to analyze")
+    """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.")
+
 
 class HypothesisInput(ResearchInput):
-    """Schema for hypothesis testing"""
-    group_col: str = Field(..., description="Categorical column defining groups")
-    value_col: str = Field(..., description="Numerical column to compare")
+    """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.")
+
 
 class GroqResearcher:
-    """Advanced AI Research Engine using Groq"""
+    """
+    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.
+    """
+
     def __init__(self, model_name="mixtral-8x7b-32768"):
         self.model_name = model_name
-        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}
+        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.
         
-        User Question: {query}
-        
-        Required Format:
-        - Executive Summary (1 paragraph)
-        - Methodology (bullet points)
-        - Key Findings (numbered list)
-        - Limitations
-        - Recommended Next Steps"""
+        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.
+
+        """
 
     def research(self, query: str, data: pd.DataFrame) -> str:
-        """Conduct academic-level analysis using Groq"""
+        """Executes in-depth research using the Groq API to produce academic-quality analyses."""
         try:
-            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
-            )
-            
+            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)
+
             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,
@@ -73,20 +89,22 @@ class GroqResearcher:
                 max_tokens=4096,
                 stream=False
             )
-            
             return completion.choices[0].message.content
-        
         except Exception as e:
-            return f"Research Error: {str(e)}"
+            return f"Research Error Encountered: {str(e)}"
+
 
 @tool(args_schema=ResearchInput)
 def advanced_eda(data_key: str) -> Dict:
-    """Comprehensive Exploratory Data Analysis with Statistical Profiling"""
+    """
+    Performs a comprehensive Exploratory Data Analysis, including statistical profiling,
+    temporal analysis of datetime columns, and detailed quality checks.
+    """
     try:
         data = st.session_state[data_key]
         analysis = {
             "dimensionality": {
-                "rows": len(data),
+                "rows": int(len(data)),  # Ensure rows are an integer
                 "columns": list(data.columns),
                 "memory_usage": f"{data.memory_usage().sum() / 1e6:.2f} MB"
             },
@@ -94,112 +112,147 @@ def advanced_eda(data_key: str) -> Dict:
             "temporal_analysis": {
                 "date_ranges": {
                     col: {
-                        "min": data[col].min(),
-                        "max": data[col].max()
+                        "min": str(data[col].min()),  # Ensure date is a string
+                        "max": str(data[col].max())  # Ensure date is a string
                     } for col in data.select_dtypes(include='datetime').columns
                 }
             },
             "data_quality": {
                 "missing_values": data.isnull().sum().to_dict(),
-                "duplicates": data.duplicated().sum(),
+                "duplicates": int(data.duplicated().sum()),  # Ensure duplicates are an integer
                 "cardinality": {
-                    col: data[col].nunique() for col in data.columns
+                    col: int(data[col].nunique()) for col in data.columns  # Ensure cardinality is integer
                 }
             }
         }
         return analysis
     except Exception as e:
-        return {"error": f"EDA Failed: {str(e)}"}
+        return {"error": f"Advanced EDA Failed: {str(e)}"}
 
 @tool(args_schema=ResearchInput)
 def visualize_distributions(data_key: str, columns: List[str]) -> str:
-    """Generate publication-quality distribution visualizations"""
+    """
+    Generates high-quality, publication-ready distribution visualizations (histograms with KDE)
+    for selected numerical columns, and returns the image as a base64 encoded string.
+    """
     try:
         data = st.session_state[data_key]
-        plt.figure(figsize=(12, 6))
+        plt.figure(figsize=(15, 7))  # Adjusted figure size for better readability
         for i, col in enumerate(columns, 1):
             plt.subplot(1, len(columns), i)
-            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()
+            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
         
         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"Visualization Error: {str(e)}"
+        return f"Distribution Visualization Error: {str(e)}"
+
 
 @tool(args_schema=TemporalAnalysisInput)
 def temporal_analysis(data_key: str, time_col: str, value_col: str) -> Dict:
-    """Time Series Decomposition and Trend Analysis"""
+    """
+    Performs a sophisticated time series analysis, including decomposition and trend assessment,
+    providing both statistical insights and a visual representation.
+    """
     try:
         data = st.session_state[data_key]
-        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))
+        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))
         decomposition.plot()
         plt.tight_layout()
-        
+
         buf = io.BytesIO()
-        plt.savefig(buf, format='png')
+        plt.savefig(buf, format='png', dpi=300)  # Increased dpi for higher resolution
         plt.close()
         plot_data = base64.b64encode(buf.getvalue()).decode()
-        
+         
+        adf_result = adfuller(ts_data)
+        stationarity_p_value = adf_result[1]
+
         return {
             "trend_statistics": {
-                "stationarity": adfuller(ts_data)[1],
-                "seasonality_strength": max(decomposition.seasonal)
+                "stationarity": stationarity_p_value,
+                "stationarity_interpretation": interpret_p_value(stationarity_p_value),
+                "seasonality_strength": max(decomposition.seasonal) if hasattr(decomposition, 'seasonal') else None
             },
-            "visualization": plot_data
+            "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,
+            }
         }
     except Exception as e:
-        return {"error": f"Temporal Analysis Failed: {str(e)}"}
+        return {"error": f"Temporal Analysis Failure: {str(e)}"}
 
 @tool(args_schema=HypothesisInput)
 def hypothesis_testing(data_key: str, group_col: str, value_col: str) -> Dict:
-    """Statistical Hypothesis Testing with Automated Assumption Checking"""
+    """
+    Conducts statistical hypothesis testing, providing detailed test results, effect size measures,
+    and interpretations for both t-tests and ANOVAs.
+    """
     try:
         data = st.session_state[data_key]
         groups = data[group_col].unique()
         
         if len(groups) < 2:
-            return {"error": "Insufficient groups for comparison"}
+            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."}
             
         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": 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)
+            "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
         }
     except Exception as e:
         return {"error": f"Hypothesis Testing Failed: {str(e)}"}
 
 def interpret_p_value(p: float) -> str:
-    """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"
+    """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)."
 
 def main():
     st.set_page_config(page_title="AI Research Lab", layout="wide")
@@ -217,9 +270,12 @@ def main():
         uploaded_file = st.file_uploader("Upload research dataset", type=["csv", "parquet"])
         if uploaded_file:
             with st.spinner("Initializing dataset..."):
-                st.session_state.data = pd.read_csv(uploaded_file)
-                st.success(f"Loaded {len(st.session_state.data):,} research observations")
-    
+                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}")
+        
     # Main research interface
     if st.session_state.data is not None:
         col1, col2 = st.columns([1, 3])
@@ -230,10 +286,10 @@ def main():
                 "Variables": list(st.session_state.data.columns),
                 "Time Range": {
                     col: {
-                        "min": st.session_state.data[col].min(),
-                        "max": st.session_state.data[col].max()
+                        "min": str(st.session_state.data[col].min()),
+                        "max": str(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"
             })
         
@@ -254,35 +310,42 @@ def main():
                     st.json(eda_result)
                 
                 elif analysis_type == "Temporal Pattern Analysis":
-                    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)
+                    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)
                 
                 elif analysis_type == "Comparative Statistics":
-                    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)
+                    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)
                 
                 elif analysis_type == "Distribution Analysis":
                     num_cols = st.session_state.data.select_dtypes(include=np.number).columns.tolist()
@@ -292,7 +355,8 @@ def main():
                             "data_key": "data",
                             "columns": selected_cols
                         })
-                        st.image(f"data:image/png;base64,{img_data}")
+                        st.image(f"data:image/png;base64,{img_data}",
+                                 use_column_width=True)
             
             with research_tab:
                 research_query = st.text_area("Enter Research Question:", height=150,