Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| import pandas as pd | |
| import numpy as np | |
| import requests | |
| import json | |
| import time | |
| import matplotlib.pyplot as plt | |
| import seaborn as sns | |
| # --- CONFIG --- | |
| # Note: GEMINI_API_KEY is retrieved from environment variables/secrets. | |
| GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") | |
| if not GEMINI_API_KEY: | |
| st.error("❌ Missing Gemini API key. Add it as a secret: GEMINI_API_KEY") | |
| st.stop() | |
| # Define API endpoints and models | |
| GEMINI_BASE = "https://generativelanguage.googleapis.com/v1beta" | |
| # Using the correct model for structured output | |
| CHAT_MODEL = "gemini-2.5-flash-preview-09-2025" | |
| EMBED_MODEL = "models/embedding-001" | |
| # Define the JSON schema for structured output | |
| ANALYSIS_SCHEMA = { | |
| "type": "OBJECT", | |
| "properties": { | |
| "reasoning": { | |
| "type": "STRING", | |
| "description": "A detailed natural language explanation of the analysis, including key findings and context." | |
| }, | |
| "code": { | |
| "type": "STRING", | |
| "description": "The complete, runnable Python code using pandas (df) and streamlit (st). Use st.pyplot() for plots, and st.dataframe() for resulting DataFrames. If no code is needed, this should be an empty string." | |
| } | |
| } | |
| } | |
| SYSTEM_INSTRUCTION = ( | |
| "You are a world-class Data Analyst Agent. Your task is to analyze the provided DataFrame ('df') " | |
| "based on the user's question. You MUST respond with a single JSON object conforming to the provided schema. " | |
| "1. **Reasoning:** Explain your plan, the steps taken, and the insights derived from the data. Format this in Markdown. " | |
| "2. **Code:** If the question requires calculation, aggregation, or visualization, you MUST generate Python code to execute against the 'df' DataFrame. " | |
| " - The DataFrame is already loaded as a variable named 'df'. Do NOT redefine it. " | |
| " - Use Streamlit functions for simple outputs: `st.dataframe(...)`, `st.bar_chart()`, `st.line_chart()`. " | |
| " - For **ALL** custom, complex plots, you MUST follow this strict Matplotlib sequence: **Start with `plt.figure()`, use `plt.` or `sns.` commands for plotting, and explicitly end with `st.pyplot(plt)`** to display the output. " | |
| " - **CRITICAL GUARDRAIL:** When generating code that uses logical conditions (e.g., in `if` statements or for complex filters) on Pandas Series or NumPy arrays, you **MUST** resolve ambiguity by using `.any()` or `.all()`. Do NOT compare a series directly to a single boolean value." | |
| " - Ensure the code is self-contained and ready to execute." | |
| ) | |
| # --- Helper Functions --- | |
| def chat_with_gemini(prompt, context): | |
| """Sends a prompt and data context to the Gemini model for structured analysis (reasoning + code).""" | |
| # Correctly prepend 'models/' to the model name in the URL path | |
| url = f"{GEMINI_BASE}/models/{CHAT_MODEL}:generateContent?key={GEMINI_API_KEY}" | |
| # Construct the full prompt including the data context | |
| full_prompt = f"Data Context (DataFrame Head and Columns):\n{context}\n\nUser Question: {prompt}" | |
| payload = { | |
| "contents": [ | |
| {"parts": [{"text": full_prompt}]} | |
| ], | |
| "systemInstruction": {"parts": [{"text": SYSTEM_INSTRUCTION}]}, | |
| "generationConfig": { | |
| "responseMimeType": "application/json", | |
| "responseSchema": ANALYSIS_SCHEMA | |
| } | |
| } | |
| max_retries = 5 | |
| delay = 1 | |
| for attempt in range(max_retries): | |
| try: | |
| r = requests.post(url, headers={'Content-Type': 'application/json'}, data=json.dumps(payload)) | |
| r.raise_for_status() | |
| data = r.json() | |
| json_str = data["candidates"][0]["content"]["parts"][0]["text"] | |
| return json.loads(json_str) | |
| except requests.exceptions.RequestException as e: | |
| if attempt < max_retries - 1: | |
| time.sleep(delay) | |
| delay *= 2 | |
| else: | |
| st.error(f"API Request Failed: {e}") | |
| raise e | |
| except Exception as e: | |
| st.error(f"Failed to parse model response or execute operation: {e}") | |
| raise e | |
| # --- UI --- | |
| st.title("✨Data Analyst Agent (Code Execution Enabled)") | |
| st.write("Upload a CSV file and ask natural language questions. The agent now generates and executes Python code to provide precise data analysis and visualizations.") | |
| # State variable to hold the DataFrame, initialized once | |
| if 'df' not in st.session_state: | |
| st.session_state.df = pd.DataFrame() | |
| uploaded = st.file_uploader("Upload CSV", type=["csv"]) | |
| if uploaded: | |
| # Use st.cache_data to avoid reloading the file multiple times | |
| def load_data(file): | |
| try: | |
| return pd.read_csv(file) | |
| except Exception as e: | |
| st.error(f"Failed to load CSV: {e}") | |
| return pd.DataFrame() | |
| st.session_state.df = load_data(uploaded) | |
| if not st.session_state.df.empty: | |
| st.subheader("Data Preview (First 5 Rows)") | |
| st.dataframe(st.session_state.df.head()) | |
| question = st.text_area("Ask a complex question or request a visualization (e.g., 'Show the average of the 'Sales' column', 'Plot the distribution of 'Age'):") | |
| if st.button("Analyze & Execute") and question: | |
| df = st.session_state.df # Local variable for code execution context | |
| # Summarize dataset for context sent to the LLM | |
| context = f"Dataset Columns: {', '.join(df.columns.astype(str))}\n\nFirst 5 rows of data:\n{df.head(5).to_string(index=False)}" | |
| st.markdown("---") | |
| st.subheader("🤖 Analysis Steps") | |
| with st.spinner("1. Generating analysis plan and code..."): | |
| try: | |
| # 1. Get structured response from LLM | |
| analysis_result = chat_with_gemini(question, context) | |
| reasoning = analysis_result.get('reasoning', "No reasoning provided.") | |
| code = analysis_result.get('code', "") | |
| st.markdown("#### 💬 Reasoning:") | |
| st.markdown(reasoning) | |
| st.markdown("#### 🐍 Generated Code:") | |
| st.code(code, language='python') | |
| except Exception as e: | |
| st.error(f"Step 1 Failed (LLM Interaction): {e}") | |
| reasoning = "" | |
| code = "" | |
| if code: | |
| with st.spinner("2. Executing code and generating output..."): | |
| try: | |
| # 2. Execute the generated Python code safely | |
| # IMPORTANT: Create a local scope with necessary variables | |
| local_scope = { | |
| 'df': df, | |
| 'st': st, | |
| 'pd': pd, | |
| 'np': np, | |
| 'plt': plt, | |
| 'sns': sns, | |
| } | |
| # Append a neutral statement to the code to prevent implicit Streamlit display of the last value | |
| final_code = code + "\nNone" | |
| # Executing the code within the local scope | |
| exec(final_code, globals(), local_scope) | |
| # FIX: Explicitly close all Matplotlib figures to prevent cross-run contamination | |
| plt.close('all') | |
| st.success("Code execution complete. Results are displayed above.") | |
| except Exception as e: | |
| st.error(f"Step 2 Failed (Code Execution Error): The agent generated invalid code. Check the console for full traceback.") | |
| st.exception(e) | |
| else: | |
| st.info("No code was generated, as the question was purely informational.") | |
| else: | |
| st.info("The uploaded CSV file appears to be empty.") | |
| else: | |
| st.info("👆 Upload a CSV file to begin the full analysis experience.") |