import streamlit as st import pandas as pd import plotly.express as px from pandasai import Agent 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 re # Set title st.title("Data Analyzer") # API keys 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.warning("API keys for OpenAI or PandasAI are missing. Ensure both keys are set in environment variables.") # Add session reset button #if st.button("Reset Session"): #for key in list(st.session_state.keys()): #del st.session_state[key] #st.experimental_rerun() # Function to validate and clean dataset def validate_and_clean_dataset(df): # Rename columns for consistency df.columns = [col.strip().lower().replace(" ", "_") for col in df.columns] # Check for missing values if df.isnull().values.any(): st.warning("Dataset contains missing values. Consider cleaning the data.") return df # Function to load datasets into session def load_dataset_into_session(): input_option = st.radio( "Select Dataset Input:", ["Use Repo Directory Dataset", "Use Hugging Face Dataset", "Upload CSV File"], ) # Option 1: Load dataset from the repo directory if input_option == "Use Repo Directory Dataset": file_path = "./source/test.csv" if st.button("Load Dataset"): try: st.session_state.df = pd.read_csv(file_path) st.session_state.df = validate_and_clean_dataset(st.session_state.df) st.success(f"File loaded successfully from '{file_path}'!") except Exception as e: st.error(f"Error loading dataset from the repo directory: {e}") # Option 2: Load dataset from Hugging Face 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: from datasets import load_dataset dataset = load_dataset(dataset_name, split="train", trust_remote_code=True) if hasattr(dataset, "to_pandas"): st.session_state.df = dataset.to_pandas() else: st.session_state.df = pd.DataFrame(dataset) st.session_state.df = validate_and_clean_dataset(st.session_state.df) st.success(f"Hugging Face Dataset '{dataset_name}' loaded successfully!") except Exception as e: st.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 a CSV File:", type=["csv"]) if uploaded_file: try: st.session_state.df = pd.read_csv(uploaded_file) st.session_state.df = validate_and_clean_dataset(st.session_state.df) st.success("File uploaded successfully!") except Exception as e: st.error(f"Error reading uploaded file: {e}") load_dataset_into_session() # Check if the dataset and API keys are loaded if "df" in st.session_state and api_key and pandasai_api_key: # Set API keys os.environ["OPENAI_API_KEY"] = api_key os.environ["PANDASAI_API_KEY"] = pandasai_api_key df = st.session_state.df st.write("Dataset Preview:") st.write(df.head()) # Ensure the dataset preview is displayed only once # Set up PandasAI Agent try: agent = Agent(df) st.info("PandasAI Agent initialized successfully.") except Exception as e: st.error(f"Error initializing PandasAI Agent: {str(e)}") # Convert dataframe into documents try: documents = [ Document( page_content=", ".join([f"{col}: {row[col]}" for col in df.columns]), metadata={"index": index} ) for index, row in df.iterrows() ] st.info("Documents created successfully for RAG.") except Exception as e: st.error(f"Error creating documents for RAG: {str(e)}") # Set up RAG try: 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 ) st.info("RAG setup completed successfully.") except Exception as e: st.error(f"Error setting up RAG: {str(e)}") # Create tabs tab1, tab2, tab3 = st.tabs(["PandasAI Analysis", "RAG Q&A", "Data Visualization"]) with tab1: st.subheader("Data Analysis with PandasAI") pandas_question = st.text_input("Ask a question about the dataset (PandasAI):") if pandas_question: try: result = agent.chat(pandas_question) st.write("PandasAI Answer:", result) if hasattr(agent, "last_output"): st.write("PandasAI Intermediate Output:", agent.last_output) except Exception as e: st.error(f"PandasAI encountered an error: {str(e)}") # Fallback: Direct pandas filtering if "patent_number" in pandas_question.lower() and "decision" in pandas_question.lower(): try: match = re.search(r'\d{7,}', pandas_question) if match: patent_number = match.group() decision = df.loc[df['patent_number'] == int(patent_number), 'decision'] st.write(f"Fallback Answer: The decision for patent {patent_number} is '{decision.iloc[0]}'.") else: st.write("Could not extract patent number from the query.") except Exception as fallback_error: st.error(f"Fallback processing failed: {fallback_error}") with tab2: st.subheader("Q&A with RAG") rag_question = st.text_input("Ask a question about the dataset (RAG):") if rag_question: try: result = qa_chain.run(rag_question) st.write("RAG Answer:", result) except Exception as e: st.error(f"RAG encountered an error: {str(e)}") with tab3: st.subheader("Data Visualization") viz_question = st.text_input("What kind of graph would you like? (e.g., 'Show a scatter plot of salary vs experience')") if viz_question: try: result = agent.chat(viz_question) code_pattern = r'```python\n(.*?)\n```' code_match = re.search(code_pattern, result, re.DOTALL) if code_match: viz_code = code_match.group(1) exec(viz_code) else: st.write("Unable to generate the graph. Showing fallback example.") fig = px.scatter(df, x=df.columns[0], y=df.columns[1]) st.plotly_chart(fig) except Exception as e: st.error(f"An error occurred during visualization: {str(e)}") else: if not api_key: st.warning("Please set the OpenAI API key in environment variables.") if not pandasai_api_key: st.warning("Please set the PandasAI API key in environment variables.")