Copied from other repo

.env

@@ -0,0 +1,2 @@
+DATABASE_URL=postgres://default:lyzegA2r0ESO@ep-dawn-fire-a1i3ytre-pooler.ap-southeast-1.aws.neon.tech/verceldb
+VANNA_API_KEY=370dd4dc5e75478f88c71f4db5cca094

app.py

@@ -0,0 +1,166 @@
+import streamlit as st
+import pandas as pd
+import plotly.express as px
+from langchain_community.llms import OpenAI
+from langchain.agents.agent_types import AgentType
+from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
+import textwrap
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from functools import partial
+import time
+
+# Initialize session state
+if 'step' not in st.session_state:
+    st.session_state.step = 1
+if 'dataframes' not in st.session_state:
+    st.session_state.dataframes = {}
+if 'chat_history' not in st.session_state:
+    st.session_state.chat_history = []
+if 'cleaning_operations' not in st.session_state:
+    st.session_state.cleaning_operations = {}
+
+def main():
+    st.title("Data Analysis Chat App")
+
+    if st.session_state.step == 1:
+        step_1_upload_and_analyze()
+    elif st.session_state.step == 2:
+        step_2_clean_data()
+    elif st.session_state.step == 3:
+        step_3_chat_with_data()
+
+def step_1_upload_and_analyze():
+    st.subheader("Step 1: Upload and Analyze Data")
+
+    uploaded_files = st.file_uploader("Upload CSV files", type="csv", accept_multiple_files=True)
+    if uploaded_files:
+        for file in uploaded_files:
+            df = pd.read_csv(file)
+            st.session_state.dataframes[file.name] = df
+            st.success(f"Uploaded: {file.name}")
+
+        if st.button("Analyze Data"):
+            for name, df in st.session_state.dataframes.items():
+                st.write(f"Analysis for {name}:")
+                st.write(f"Shape: {df.shape}")
+                st.write("Columns:")
+                st.write(df.columns.tolist())
+                st.write("Preview:")
+                st.write(df.head())
+                st.write("---")
+
+        if st.button("Proceed to Data Cleaning"):
+            st.session_state.step = 2
+
+def step_2_clean_data():
+    st.subheader("Step 2: Clean Data")
+
+    llm = OpenAI(temperature=0)
+    
+    for name, df in st.session_state.dataframes.items():
+        st.write(f"Cleaning recommendations for {name}:")
+        
+        # Create a summary of the dataframe
+        summary = f"Dataframe '{name}' summary:\n"
+        summary += f"- Shape: {df.shape}\n"
+        summary += f"- Columns: {', '.join(df.columns)}\n"
+        summary += "- Data types:\n"
+        for col, dtype in df.dtypes.items():
+            summary += f"  - {col}: {dtype}\n"
+        summary += "- Sample data (first 5 rows):\n"
+        summary += df.head().to_string()
+
+        # Split the summary into smaller chunks
+        chunk_size = 1500  # Reduced chunk size
+        chunks = textwrap.wrap(summary, chunk_size)
+
+        cleaning_recommendations = []
+        with st.spinner("Analyzing data and generating recommendations..."):
+            for i, chunk in enumerate(chunks):
+                chunk_result = analyze_chunk(llm, df, chunk)
+                cleaning_recommendations.append(chunk_result)
+
+        # Combine all recommendations
+        full_recommendations = "\n".join(cleaning_recommendations)
+        st.write(full_recommendations)
+        
+        # Create checkboxes for cleaning operations
+        cleaning_ops = [op.strip() for op in full_recommendations.split('\n') if op.strip()]
+        st.session_state.cleaning_operations[name] = []
+        for op in cleaning_ops:
+            if st.checkbox(op, key=f"{name}_{op}"):
+                st.session_state.cleaning_operations[name].append(op)
+
+    if st.button("Apply Cleaning and Proceed to Chat"):
+        for name, ops in st.session_state.cleaning_operations.items():
+            df = st.session_state.dataframes[name]
+            for op in ops:
+                # Here you would implement the actual cleaning operations
+                # For now, we'll just print what would be done
+                st.write(f"Applying to {name}: {op}")
+        
+        st.session_state.step = 3
+        st.success("Cleaning operations applied. Proceeding to chat interface.")
+        st.button("Go to Chat Interface")
+
+    if st.button("Back to Data Upload"):
+        st.session_state.step = 1
+        st.experimental_rerun()
+
+def step_3_chat_with_data():
+    st.subheader("Step 3: Chat with your data")
+
+    user_input = st.text_input("Ask a question about your data:")
+    if user_input:
+        response = process_user_input(user_input)
+        st.session_state.chat_history.append(("User", user_input))
+        st.session_state.chat_history.append(("AI", response))
+
+    for role, message in st.session_state.chat_history:
+        if role == "User":
+            st.text_area("You:", value=message, height=50, disabled=True)
+        else:
+            st.text_area("AI:", value=message, height=100, disabled=True)
+
+def process_user_input(user_input):
+    llm = OpenAI(temperature=0)
+    combined_df = pd.concat([df.assign(source=name) for name, df in st.session_state.dataframes.items()], ignore_index=True)
+
+    df_summary = "Available data:\n"
+    for name, df in st.session_state.dataframes.items():
+        df_summary += f"- {name}: {len(df)} rows, {len(df.columns)} columns\n"
+        df_summary += f"  Columns: {', '.join(df.columns)}\n\n"
+
+    agent = create_pandas_dataframe_agent(
+        llm,
+        combined_df,
+        verbose=True,
+        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        allow_dangerous_code=True
+    )
+
+    full_input = f"{df_summary}\nThe data from all files has been combined into a single DataFrame with an additional 'source' column indicating the original file.\n\nUser question: {user_input}"
+
+    response = agent.run(full_input)
+    return response
+
+def analyze_chunk(llm, df, chunk, timeout=30):
+    agent = create_pandas_dataframe_agent(
+        llm,
+        df,
+        verbose=True,
+        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        allow_dangerous_code=True
+    )
+    
+    prompt = f"Analyze this part of the dataframe summary and suggest up to 3 specific cleaning operations. Focus on identifying missing values, outliers, and inconsistent data formats.\n\n{chunk}"
+    
+    try:
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(agent.run, prompt)
+            return future.result(timeout=timeout)
+    except Exception as e:
+        return f"Analysis timed out or encountered an error: {str(e)}"
+
+if __name__ == "__main__":
+    main()

explore.py

@@ -0,0 +1,83 @@
+import streamlit as st
+import os
+from vanna import VannaBase
+import pandas as pd
+from dotenv import load_dotenv
+from sqlalchemy import create_engine
+from sqlalchemy.exc import SQLAlchemyError
+
+# Load environment variables
+load_dotenv()
+
+# Initialize Vanna AI
+vanna_api_key = os.getenv("VANNA_API_KEY")
+if not vanna_api_key:
+    st.error("VANNA_API_KEY is not set in the environment variables. Please set it and restart the application.")
+    st.stop()
+
+vn = VannaBase(api_key=vanna_api_key)
+
+# Check if DATABASE_URL is set
+database_url = os.getenv("DATABASE_URL")
+if not database_url:
+    st.error("DATABASE_URL is not set in the environment variables. Please set it and restart the application.")
+    st.stop()
+
+# Try to connect to the database
+try:
+    engine = create_engine(database_url)
+    with engine.connect() as connection:
+        st.success("Successfully connected to the database.")
+    vn.connect_to_postgres(database_url)
+except SQLAlchemyError as e:
+    st.error(f"Failed to connect to the database: {str(e)}")
+    st.stop()
+
+st.title("Data Explorer")
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Chat input
+if prompt := st.chat_input("Ask about your data"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(prompt)
+    
+    try:
+        # Generate SQL query
+        sql_query = vn.generate_sql(prompt)
+        
+        # Execute SQL query and get results
+        df = vn.run_sql(sql_query)
+        
+        # Display assistant response
+        with st.chat_message("assistant"):
+            st.markdown(f"Here's the SQL query I generated:\n```sql\n{sql_query}\n```")
+            st.markdown("And here are the results:")
+            st.dataframe(df)
+        
+        # Add assistant message to chat history
+        st.session_state.messages.append({
+            "role": "assistant",
+            "content": f"Here's the SQL query I generated:\n```sql\n{sql_query}\n```\n\nAnd here are the results:\n{df.to_markdown()}"
+        })
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")
+
+# Sidebar with additional information
+st.sidebar.header("About")
+st.sidebar.info(
+    "This is a data exploration tool using Streamlit and Vanna AI. "
+    "Ask questions about your data in natural language, and the app will "
+    "generate SQL queries and display the results."
+)

persistence.py

@@ -0,0 +1,80 @@
+import os
+import pandas as pd
+import time
+from sqlalchemy import create_engine, Column, String, Integer, Float, DateTime, inspect, MetaData
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.exc import SQLAlchemyError
+
+
+DATABASE_URL = os.environ.get('DATABASE_URL')
+engine = create_engine(DATABASE_URL)
+
+
+def get_df_from_csv(csv_file_path):
+    df = pd.read_csv(csv_file_path)
+    return df
+
+def get_schema_from_df(df):
+    schema = pd.io.json.build_table_schema(df)
+    return schema
+
+def create_table_from_schema(table_name, schema):
+    Base = declarative_base()
+    
+    inspector = inspect(engine)
+    metadata = MetaData()
+    metadata.reflect(bind=engine)
+
+    # Check if table already exists
+    if table_name in inspector.get_table_names():
+        existing_columns = {column['name']: column['type'] for column in inspector.get_columns(table_name)}
+        new_columns = {field['name']: field['type'] for field in schema['fields']}
+        
+        if existing_columns == new_columns:
+            print(f"Table '{table_name}' with the same schema already exists. Skipping creation.")
+            return
+        else:
+            print(f"Table '{table_name}' exists but has a different schema. Creating a new table with a timestamp suffix.")
+            table_name = f"{table_name}_{int(time.time())}"
+
+    class DynamicTable(Base):
+        __tablename__ = table_name
+        
+        id = Column(Integer, primary_key=True)
+        
+        for column in schema['fields']:
+            if column['name'] != 'id':
+                if column['type'] == 'integer':
+                    locals()[column['name']] = Column(Integer)
+                elif column['type'] == 'number':
+                    locals()[column['name']] = Column(Float)
+                elif column['type'] == 'datetime':
+                    locals()[column['name']] = Column(DateTime)
+                else:
+                    locals()[column['name']] = Column(String)
+    
+    try:
+        Base.metadata.create_all(engine)
+        print(f"Table '{table_name}' created successfully.")
+    except SQLAlchemyError as e:
+        print(f"Error creating table: {str(e)}")
+
+
+def save_data_to_table(table_name, df):
+    try:
+        df.to_sql(table_name, engine)
+    except SQLAlchemyError as e:
+        print(f"Error saving data to table: {str(e)}")
+    
+
+
+if __name__ == "__main__":
+    filename = 'data.csv'
+    df = get_df_from_csv(filename)
+    schema = get_schema_from_df(df)    
+    table_name = filename.split('.')[0]
+    
+    create_table_from_schema(table_name, schema)
+    save_data_to_table(table_name, df)
+
+

requirements.txt

@@ -0,0 +1,19 @@
+plotly==5.24.1
+langchain==0.3.4
+langchain-community==0.3.3
+langchain-core==0.3.12
+langchain-experimental==0.3.2
+langchain-openai==0.2.3
+langchain-text-splitters==0.3.0
+tabulate==0.9.0
+vanna==0.7.3
+psycopg2-binary
+psycopg2
+streamlit==1.31.0
+pandas==2.2.0
+python-dotenv==1.0.0
+sqlalchemy==2.0.25
+
+
+
+>>>>>>> 8798f85 (add deps)

tryvanna.py

@@ -0,0 +1,23 @@
+import os
+
+from vanna.remote import VannaDefault
+vn = VannaDefault(model='gpt-3.5-turbo', api_key=os.getenv("VANNA_API_KEY"))
+
+# # vn.connect_to_postgres(os.getenv("DATABASE_URL"))
+#  export DATABASE_NAME=verceldb
+# ashwin@MacBook-Air-6 DDD % export DATABASE_USER=default
+# ashwin@MacBook-Air-6 DDD % export DATABASE_PASSWORD=lyzegA2r0ESO
+# ashwin@MacBook-Air-6 DDD % export DATABASE_HOST="ep-dawn-fire-a1i3ytre-pooler.ap-southeast-1.aws.neon.tech" 
+
+db_host = os.getenv("DATABASE_HOST")
+db_name = os.getenv("DATABASE_NAME")
+db_user = os.getenv("DATABASE_USER")
+db_password = os.getenv("DATABASE_PASSWORD")
+db_port = 5432
+
+vn.connect_to_postgres(host=db_host, dbname=db_name, user=db_user, password=db_password, port=db_port)
+
+vn.ask('What are the top 10 artists by sales?')
+
+from vanna.flask import VannaFlaskApp
+VannaFlaskApp(vn).run()