import streamlit as st
import pandas as pd
import logging
from IPython.display import Markdown, display
from llama_index.core.query_engine import PandasQueryEngine
from llama_index.llms.anthropic import Anthropic

logging.basicConfig(level=logging.INFO)

def main():
    st.title('CSV Data Query Tool')

    # Description
    st.write("""
    This app helps you query tabular data using natural language. 
    Upload a CSV file, enter your ANTHROPIC_API_KEY, ask a question about the data, and get a response!
    """)

    # Upload CSV file
    uploaded_file = st.file_uploader("Upload CSV file", type=['csv'])
    if uploaded_file is not None:
        df = pd.read_csv(uploaded_file)

        # Input API Key
        api_key = st.text_input("Enter your ANTHROPIC_API_KEY:")

        # Initialize Anthropic
        if api_key:
            llm = Anthropic(temperature=0.0, model='claude-3-opus-20240229', api_key=api_key)
            query_engine = PandasQueryEngine(df=df, llm=llm, verbose=True)

            # Input user question
            question = st.text_input("Ask a question about the data:")
            if st.button("Submit"):
                if question:
                    response = query_engine.query(question)
                    st.markdown(f"**Response:** {response}")

if __name__ == "__main__":
    main()