import streamlit as st import pandas as pd import plotly.express as px from datasets import load_dataset from pandasai import Agent from pandasai.llm.openai import OpenAI from langchain_community.embeddings.openai import OpenAIEmbeddings from langchain_community.vectorstores import FAISS from langchain_openai import ChatOpenAI from langchain.chains import RetrievalQA from langchain.schema import Document import os import logging # Configure logging logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) # Set the title of the app st.title("Data Analyzer") # Fetch API keys from environment variables api_key = os.getenv("OPENAI_API_KEY") pandasai_api_key = os.getenv("PANDASAI_API_KEY") if not api_key or not pandasai_api_key: st.error( "API keys not found in the environment. Please set the 'OPENAI_API_KEY' and 'PANDASAI_API_KEY' environment variables." ) logger.error("API keys not found. Ensure they are set in the environment variables.") else: def load_dataset_into_session(): """Function to load a dataset into the session.""" input_option = st.radio("Select Dataset Input:", ["Use Repo Dataset", "Use Hugging Face Dataset", "Upload CSV File"]) # Option 1: Use Repo Dataset if input_option == "Use Repo Dataset": file_path = "./source/test.csv" if st.button("Load Repo Dataset"): try: st.session_state.df = pd.read_csv(file_path) st.success(f"File loaded successfully from '{file_path}'!") st.dataframe(st.session_state.df.head(10)) except Exception as e: st.error(f"Error reading file from path: {e}") logger.error(f"Error reading file from path: {e}") # Option 2: Use Hugging Face Dataset elif input_option == "Use Hugging Face Dataset": dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd") if st.button("Load Hugging Face Dataset"): try: # Load Hugging Face dataset dataset = load_dataset(dataset_name, split="train", trust_remote_code=True) st.session_state.df = pd.DataFrame(dataset) st.success(f"Dataset '{dataset_name}' loaded successfully!") st.dataframe(st.session_state.df.head(10)) except Exception as e: st.error(f"Error loading dataset from Hugging Face: {e}") logger.error(f"Error loading Hugging Face dataset: {e}") # Option 3: Upload CSV File elif input_option == "Upload CSV File": uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"]) if uploaded_file: try: st.session_state.df = pd.read_csv(uploaded_file) st.success("File uploaded successfully!") st.dataframe(st.session_state.df.head(10)) except Exception as e: st.error(f"Error reading uploaded file: {e}") logger.error(f"Error reading uploaded file: {e}") # Initialize session state for DataFrame if "df" not in st.session_state: st.session_state.df = None # Load dataset into session load_dataset_into_session() # Proceed only if a DataFrame is loaded if st.session_state.df is not None: df = st.session_state.df try: # Initialize PandasAI Agent llm = OpenAI(api_key=pandasai_api_key, max_tokens=1500, timeout=60) agent = Agent(df, llm=llm) # Convert the DataFrame into documents for RAG documents = [ Document( page_content=", ".join([f"{col}: {row[col]}" for col in df.columns if pd.notnull(row[col])]), metadata={"index": index} ) for index, row in df.iterrows() ] logger.info(f"{len(documents)} documents created for RAG.") # Set up RAG embeddings = OpenAIEmbeddings() vectorstore = FAISS.from_documents(documents, embeddings) retriever = vectorstore.as_retriever() qa_chain = RetrievalQA.from_chain_type( llm=ChatOpenAI(), chain_type="stuff", retriever=retriever ) # Create tabs tab1, tab2, tab3 = st.tabs(["PandasAI Analysis", "RAG Q&A", "Data Visualization"]) # Tab 1: PandasAI Analysis with tab1: st.header("Data Analysis using PandasAI") pandas_question = st.text_input("Ask a question about the data (PandasAI):") if pandas_question: try: result = agent.chat(pandas_question) if result: st.write("PandasAI Answer:", result) else: st.warning("PandasAI returned no result. Please try another question.") except Exception as e: st.error(f"Error from PandasAI: {e}") logger.error(f"PandasAI error: {e}") # Tab 2: RAG Q&A with tab2: st.header("Question Answering using RAG") rag_question = st.text_input("Ask a question about the data (RAG):") if rag_question: try: result = qa_chain.run(rag_question) st.write("RAG Answer:", result) except Exception as e: st.error(f"Error from RAG Q&A: {e}") logger.error(f"RAG error: {e}") # Tab 3: Data Visualization with tab3: st.header("Data Visualization") viz_question = st.text_input("What kind of graph would you like to create? (e.g., 'Show a scatter plot of salary vs experience')") if viz_question: try: result = agent.chat(viz_question) # Extract Python code for visualization code_pattern = r'```python\n(.*?)\n```' code_match = re.search(code_pattern, result, re.DOTALL) if code_match: viz_code = code_match.group(1) logger.debug(f"Extracted visualization code: {viz_code}") # Safeguard: Modify and validate code for Plotly viz_code = viz_code.replace('plt.', 'px.') exec(viz_code) # Execute the visualization code st.plotly_chart(fig) else: st.warning("Unable to generate a graph. Please try a different query.") logger.warning("No valid visualization code found in PandasAI response.") except Exception as e: st.error(f"An error occurred: {e}") logger.error(f"Visualization error: {e}") except Exception as e: st.error(f"An error occurred while processing the dataset: {e}") logger.error(f"Dataset processing error: {e}")