Update app.py

app.py

@@ -69,33 +69,42 @@ def standardize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
     return df
 
 def analyze_columns(df: pd.DataFrame, filename: str) -> dict:
-    """Analyze DataFrame columns using Gemini AI with improved error handling."""
+    """Analyze DataFrame columns using Gemini AI with improved error handling and prompt."""
     try:
         display_df = df.head(5).copy()
         for col in display_df.columns:
             display_df[col] = display_df[col].astype(str)
         sample_csv = display_df.to_csv(index=False)
         prompt = f"""
-        Analyze this CSV data and provide analysis in JSON format.
+        Analyze this CSV data, which represents employee payroll information, and provide analysis in JSON format.
+
         Filename: {filename}
-        Sample data:
+        Sample data (first 5 rows):
         {sample_csv}
-        Respond with only a valid JSON object in this format:
+
+        In the context of merging datasets, "key columns" are columns that uniquely identify records and are essential for joining this data with other datasets. For payroll data, key columns are typically employee identifiers such as Employee ID, Taxpayer Identification Number (TIN), or Employee Name (if unique).
+
+        Please analyze the columns in the sample data and identify potential key columns that can be used to merge this dataset with other employee-related datasets.
+
+        Respond with ONLY a valid JSON object in the following format:
+
         {{
-            "subject": "Employee payroll data",
+            "subject": "Employee payroll data analysis",
             "columns": [
                 {{
                     "name": "column_name",
                     "type": "string/number/date",
-                    "description": "Brief description"
+                    "description": "Brief description of the column and its likely content."
                 }}
             ],
-            "key_columns": ["employee_id", "tin"],
-            "issues": ["Missing values in salary column"],
+            "key_columns": ["List of identified key column names.  Prioritize employee identifiers like employee_id, tin, or employee_name if they appear to be unique identifiers."],
+            "issues": ["List any data quality issues found, like missing values in important columns."],
             "suggested_renames": {{
                 "old_name": "new_name"
             }}
         }}
+
+        Ensure the JSON response is valid and parsable. Focus on accurately identifying key columns relevant for merging payroll data.
         """
         response = model.generate_content(prompt)
         response_text = response.text.strip()
@@ -156,12 +165,26 @@ def merge_with_master(processed_files):
         st.warning("No master file with 'earnings' found. Using the first file as master.")
         master_file = processed_files[0]
         other_files = processed_files[1:]
+
+    if not master_file: # Handle case where no files are processed correctly
+        st.error("No master file could be determined. Merging cannot proceed.")
+        return None
+
     master_df = master_file["df"]
     master_keys = master_file["analysis"].get("key_columns", [])
     st.write(f"Using '{master_file['filename']}' as master with key columns: {master_keys}")
     merged_df = master_df
+
+    if merged_df.empty: # Check if master_df is empty. If so, no point merging.
+        st.warning(f"Master DataFrame '{master_file['filename']}' is empty. Merging will result in an empty DataFrame.")
+        return merged_df
+
     for other in other_files:
         other_df = other["df"]
+        if other_df.empty: # Check if other_df is empty before merging
+            st.warning(f"DataFrame '{other['filename']}' is empty. Skipping merge for this file.")
+            continue
+
         other_keys = other["analysis"].get("key_columns", [])
         common_keys = list(set(master_keys).intersection(set(other_keys)))
         if common_keys:
@@ -192,6 +215,9 @@ def main():
                 else:
                     df = pd.read_csv(uploaded_file)
                 if df is not None:
+                    if df.empty: # Check if dataframe is empty immediately after reading.
+                        st.warning(f"DataFrame from '{uploaded_file.name}' is empty after reading. Please check the file.")
+                        continue # Skip to next file
                     df.columns = [clean_column_name(col) for col in df.columns]
                     df = standardize_dataframe(df)
                     st.write("Initial Preview:")
@@ -207,39 +233,49 @@ def main():
                         processed_files.append(
                             {"filename": uploaded_file.name, "df": df, "analysis": analysis}
                         )
+                else:
+                    st.error(f"Could not read data from '{uploaded_file.name}'.") # Explicit error if read_excel_file returns None
+
             except Exception as e:
                 st.error(f"Error processing {uploaded_file.name}: {str(e)}")
                 continue
+
         if len(processed_files) > 1:
-            st.write("### Merging DataFrames with Earnings Schedule as Master")
-            merged_df = merge_with_master(processed_files)
-            if merged_df is not None:
-                st.write("### Preview of Merged Data")
-                st.dataframe(safe_display_df(merged_df.head()))
-                try:
-                    csv = merged_df.to_csv(index=False)
-                    st.download_button(
-                        label="Download Merged CSV",
-                        data=csv,
-                        file_name="merged_data.csv",
-                        mime="text/csv",
-                    )
-                    st.write("### Dataset Statistics")
-                    st.write(f"Total rows: {len(merged_df)}")
-                    st.write(f"Total columns: {len(merged_df.columns)}")
-                    st.write("### Data Quality Metrics")
-                    missing_df = pd.DataFrame(
-                        {
-                            "Column": merged_df.columns,
-                            "Missing Values": merged_df.isnull().sum().values,
-                            "Missing Percentage": (merged_df.isnull().sum().values / len(merged_df) * 100).round(2),
-                        }
-                    )
-                    st.dataframe(missing_df)
-                    duplicates = merged_df.duplicated().sum()
-                    st.write(f"Number of duplicate rows: {duplicates}")
-                except Exception as e:
-                    st.error(f"Error preparing download: {str(e)}")
+            if not any(file_info["df"].empty for file_info in processed_files): # Check if any processed dataframe is empty before merging
+                st.write("### Merging DataFrames with Earnings Schedule as Master")
+                merged_df = merge_with_master(processed_files)
+                if merged_df is not None and not merged_df.empty: # Check merged_df is not None and not empty before displaying.
+                    st.write("### Preview of Merged Data")
+                    st.dataframe(safe_display_df(merged_df.head()))
+                    try:
+                        csv = merged_df.to_csv(index=False)
+                        st.download_button(
+                            label="Download Merged CSV",
+                            data=csv,
+                            file_name="merged_data.csv",
+                            mime="text/csv",
+                        )
+                        st.write("### Dataset Statistics")
+                        st.write(f"Total rows: {len(merged_df)}")
+                        st.write(f"Total columns: {len(merged_df.columns)}")
+                        st.write("### Data Quality Metrics")
+                        missing_df = pd.DataFrame(
+                            {
+                                "Column": merged_df.columns,
+                                "Missing Values": merged_df.isnull().sum().values,
+                                "Missing Percentage": (merged_df.isnull().sum().values / len(merged_df) * 100).round(2),
+                                }
+                            )
+                        st.dataframe(missing_df)
+                        duplicates = merged_df.duplicated().sum()
+                        st.write(f"Number of duplicate rows: {duplicates}")
+                    except Exception as e:
+                        st.error(f"Error preparing download: {str(e)}")
+                elif merged_df is not None and merged_df.empty: # Explicitly handle empty merged dataframe case
+                    st.warning("The merged DataFrame is empty. Please check the input files and merging keys.")
+            else:
+                st.warning("One or more of the processed DataFrames is empty. Merging cannot proceed meaningfully.")
+
         else:
             st.warning("Please upload at least 2 files to merge.")