Update app.py

app.py

@@ -1,193 +1,65 @@
-import os
-import json
-import re
-import sys
-import io
-import contextlib
-import warnings
-from typing import Optional, List, Any, Tuple
-from PIL import Image
-import streamlit as st
-import pandas as pd
-import base64
-from io import BytesIO
-from together import Together
-from e2b_code_interpreter import Sandbox
-
-warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")
-
-pattern = re.compile(r"```python\n(.*?)\n```", re.DOTALL)
-
-def code_interpret(e2b_code_interpreter: Sandbox, code: str) -> Optional[List[Any]]:
-    with st.spinner('Executing code in E2B sandbox...'):
-        stdout_capture = io.StringIO()
-        stderr_capture = io.StringIO()
-
-        with contextlib.redirect_stdout(stdout_capture), contextlib.redirect_stderr(stderr_capture):
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore")
-                exec = e2b_code_interpreter.run_code(code)
-
-        if stderr_capture.getvalue():
-            print("[Code Interpreter Warnings/Errors]", file=sys.stderr)
-            print(stderr_capture.getvalue(), file=sys.stderr)
-
-        if stdout_capture.getvalue():
-            print("[Code Interpreter Output]", file=sys.stdout)
-            print(stdout_capture.getvalue(), file=sys.stdout)
-
-        if exec.error:
-            print(f"[Code Interpreter ERROR] {exec.error}", file=sys.stderr)
-            return None
-        return exec.results
-
-def match_code_blocks(llm_response: str) -> str:
-    match = pattern.search(llm_response)
-    if match:
-        code = match.group(1)
-        return code
-    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.
-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."""
-
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_message},
-    ]
-
-    with st.spinner('Getting response from Together AI LLM model...'):
-        client = Together(api_key=st.session_state.together_api_key)
-        response = client.chat.completions.create(
-            model=st.session_state.model_name,
-            messages=messages,
-        )
-
-        response_message = response.choices[0].message
-        python_code = match_code_blocks(response_message.content)
-        
-        if python_code:
-            code_interpreter_results = code_interpret(e2b_code_interpreter, python_code)
-            return code_interpreter_results, response_message.content
-        else:
-            st.warning(f"Failed to match any Python code in model's response")
-            return None, response_message.content
-
-def upload_dataset(code_interpreter: Sandbox, uploaded_file) -> str:
-    dataset_path = f"./{uploaded_file.name}"
-    
-    try:
-        code_interpreter.files.write(dataset_path, uploaded_file)
-        return dataset_path
-    except Exception as error:
-        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!")
-
-    # Initialize session state variables
-    if 'together_api_key' not in st.session_state:
-        st.session_state.together_api_key = ''
-    if 'e2b_api_key' not in st.session_state:
-        st.session_state.e2b_api_key = ''
-    if 'model_name' not in st.session_state:
-        st.session_state.model_name = ''
-
-    # Sidebar for API keys and model configuration
-    with st.sidebar:
-        st.header("🔑 API Keys and Model Configuration")
-        st.session_state.together_api_key = st.text_input("Together AI API Key", type="password")
-        st.info("💡 Everyone gets a free $1 credit by Together AI - AI Acceleration Cloud platform")
-        st.markdown("[Get Together AI API Key](https://api.together.ai/signin)")
-        
-        st.session_state.e2b_api_key = st.text_input("Enter E2B API Key", type="password")
-        st.markdown("[Get E2B API Key](https://e2b.dev/docs/legacy/getting-started/api-key)")
-        
-        # Add model selection dropdown
-        model_options = {
-            "Meta-Llama 3.1 405B": "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
-            "DeepSeek V3": "deepseek-ai/DeepSeek-V3",
-            "Qwen 2.5 7B": "Qwen/Qwen2.5-7B-Instruct-Turbo",
-            "Meta-Llama 3.3 70B": "meta-llama/Llama-3.3-70B-Instruct-Turbo"
-        }
-        st.session_state.model_name = st.selectbox(
-            "Select Model",
-            options=list(model_options.keys()),
-            index=0  # Default to first option
-        )
-        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
-
-    with col1:
-        st.header("📂 Upload Dataset")
-        uploaded_file = st.file_uploader("Choose a CSV file", type="csv", key="file_uploader")
-        
-        if uploaded_file is not None:
-            # Display dataset with toggle
-            df = pd.read_csv(uploaded_file)
-            st.write("### Dataset Preview")
-            show_full = st.checkbox("Show full dataset")
-            if show_full:
-                st.dataframe(df)
-            else:
-                st.write("Preview (first 5 rows):")
-                st.dataframe(df.head())
-
-    with col2:
-        if uploaded_file is not None:
-            st.header("❓ Ask a Question")
-            query = st.text_area(
-                "What would you like to know about your data?",
-                "Can you compare the average cost for two people between different categories?",
-                height=100
-            )
-            
-            if st.button("Analyze", type="primary", key="analyze_button"):
-                if not st.session_state.together_api_key or not st.session_state.e2b_api_key:
-                    st.error("Please enter both API keys in the sidebar.")
-                else:
-                    with Sandbox(api_key=st.session_state.e2b_api_key) as code_interpreter:
-                        # Upload the dataset
-                        dataset_path = upload_dataset(code_interpreter, uploaded_file)
-                        
-                        # Pass dataset_path to chat_with_llm
-                        code_results, llm_response = chat_with_llm(code_interpreter, query, dataset_path)
-                        
-                        # Display LLM's text response
-                        st.header("🤖 AI Response")
-                        st.write(llm_response)
-                        
-                        # Display results/visualizations
-                        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
-                                    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.plotly_chart(result)
-                                elif isinstance(result, (pd.DataFrame, pd.Series)):
-                                    st.dataframe(result)
-                                else:
-                                    st.write(result)
-
-if __name__ == "__main__":
-    main()
+import pandas as pd
+import openai
+import streamlit as st
+import matplotlib.pyplot as plt
+
+# Analyze using OpenAI
+def get_openai_insights(api_key, prompt):
+    openai.api_key = api_key
+    response = openai.Completion.create(
+        engine="text-davinci-003",
+        prompt=prompt,
+        max_tokens=500,
+        temperature=0.5
+    )
+    return response["choices"][0]["text"].strip()
+
+# Streamlit app
+def main():
+    st.title("Excel Data Visualization with OpenAI Insights")
+
+    # Input OpenAI API Key
+    api_key = st.text_input("Enter your OpenAI API Key", type="password")
+    if not api_key:
+        st.warning("Please enter your OpenAI API key to proceed.")
+        return
+
+    # File upload
+    excel_file = st.file_uploader("Upload the Excel File", type=["xls", "xlsx"])
+
+    if excel_file:
+        # Load Excel data
+        excel_data = pd.ExcelFile(excel_file)
+        st.sidebar.header("Select a Sheet to Visualize")
+        sheet_name = st.sidebar.selectbox("Sheet Name", excel_data.sheet_names)
+
+        if sheet_name:
+            data = pd.read_excel(excel_data, sheet_name=sheet_name)
+            st.subheader(f"Data from Sheet: {sheet_name}")
+            st.dataframe(data)
+
+            # Option to generate insights using OpenAI
+            st.header("Generate AI Insights")
+            if st.button("Get Insights from OpenAI"):
+                with st.spinner("Generating insights..."):
+                    try:
+                        data_sample = data.head(5).to_csv(index=False)
+                        prompt = f"Analyze the following data and provide key insights:\n\n{data_sample}"
+                        insights = get_openai_insights(api_key, prompt)
+                        st.success("AI Insights Generated!")
+                        st.text_area("AI Insights:", insights, height=200)
+                    except openai.error.OpenAIError as e:
+                        st.error(f"Error with OpenAI API: {e}")
+
+            # Visualize numeric data
+            st.header("Visualize Data")
+            numeric_cols = data.select_dtypes(include="number").columns
+            if numeric_cols.any():
+                col_to_plot = st.selectbox("Select a Column to Plot", numeric_cols)
+                if col_to_plot:
+                    fig, ax = plt.subplots()
+                    data[col_to_plot].plot(kind="bar", ax=ax, title=f"{col_to_plot} Analysis")
+                    st.pyplot(fig)
+
+if __name__ == "__main__":
+    main()