Update app.py

app.py

@@ -49,10 +49,12 @@ def match_code_blocks(llm_response: str) -> str:
     return ""
 
 def chat_with_llm(e2b_code_interpreter: Sandbox, user_message: str, dataset_path: str) -> Tuple[Optional[List[Any]], str]:
-    # Update system prompt to include dataset path information
-    system_prompt = f"""You're a Python data scientist and data visualization expert. You are given a dataset at path '{dataset_path}' and also the user's query.
+    # Updated system prompt with Excel support
+    system_prompt = f"""You're a Python data scientist and data visualization expert. You are given a dataset at path '{dataset_path}' (could be CSV or Excel) and the user's query.
 You need to analyze the dataset and answer the user's query with a response and you run Python code to solve them.
-IMPORTANT: Always use the dataset path variable '{dataset_path}' in your code when reading the CSV file."""
+IMPORTANT: 
+- Use pd.read_csv() for .csv files and pd.read_excel() for .xlsx/.xls files
+- Always use the dataset path variable '{dataset_path}' in your code when reading the file"""
 
     messages = [
         {"role": "system", "content": system_prompt},
@@ -86,13 +88,12 @@ def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str:
         st.error(f"Error during file upload: {error}")
         raise error
 
-
 def main():
     """Main Streamlit application."""
     st.set_page_config(page_title="📊 AI Data Visualization Agent", page_icon="📊", layout="wide")
     
     st.title("📊 AI Data Visualization Agent")
-    st.write("Upload your dataset and ask questions about it!")
+    st.write("Upload your dataset (CSV or Excel) and ask questions about it!")
 
     # Initialize session state variables
     if 'together_api_key' not in st.session_state:
@@ -127,15 +128,20 @@ def main():
         st.session_state.model_name = model_options[st.session_state.model_name]
 
     # Main content layout
-    col1, col2 = st.columns([1, 2])  # Split the main content into two columns
+    col1, col2 = st.columns([1, 2])
 
     with col1:
         st.header("📂 Upload Dataset")
-        uploaded_file = st.file_uploader("Choose a CSV file", type="csv", key="file_uploader")
+        # Updated file uploader to accept both CSV and Excel
+        uploaded_file = st.file_uploader("Choose a CSV or Excel file", type=["csv", "xlsx", "xls"], key="file_uploader")
         
         if uploaded_file is not None:
-            # Display dataset with toggle
-            df = pd.read_csv(uploaded_file)
+            # Handle both CSV and Excel files
+            if uploaded_file.name.endswith(('.xlsx', '.xls')):
+                df = pd.read_excel(uploaded_file)
+            else:
+                df = pd.read_csv(uploaded_file)
+            
             st.write("### Dataset Preview")
             show_full = st.checkbox("Show full dataset")
             if show_full:
@@ -172,17 +178,14 @@ def main():
                         if code_results:
                             st.header("📊 Analysis Results")
                             for result in code_results:
-                                if hasattr(result, 'png') and result.png:  # Check if PNG data is available
-                                    # Decode the base64-encoded PNG data
+                                if hasattr(result, 'png') and result.png:
                                     png_data = base64.b64decode(result.png)
-                                    
-                                    # Convert PNG data to an image and display it
                                     image = Image.open(BytesIO(png_data))
-                                    st.image(image, caption="Generated Visualization", use_container_width=True)
-                                elif hasattr(result, 'figure'):  # For matplotlib figures
-                                    fig = result.figure  # Extract the matplotlib figure
-                                    st.pyplot(fig)  # Display using st.pyplot
-                                elif hasattr(result, 'show'):  # For plotly figures
+                                    st.image(image, caption="Generated Visualization", use_column_width=True)
+                                elif hasattr(result, 'figure'):
+                                    fig = result.figure
+                                    st.pyplot(fig)
+                                elif hasattr(result, 'show'):
                                     st.plotly_chart(result)
                                 elif isinstance(result, (pd.DataFrame, pd.Series)):
                                     st.dataframe(result)