pratham0011/QueryMate_Text-to-SQL-CSV

1import streamlit as st2import requests3import pandas as pd4 5st.set_page_config(page_title="QueryMate: Text to SQL & CSV")6 7st.markdown("# QueryMate: Text to SQL & CSV 💬🗄️")8st.markdown('''Welcome to QueryMate, your friendly assistant for converting natural language queries into SQL statements and CSV outputs!9               Let's get started with your data queries!''')10 11# Initialize chat history in session state if it doesn't exist12if 'chat_history' not in st.session_state:13    st.session_state.chat_history = []14 15# Data source selection16data_source = st.radio("Select Data Source:", ('SQL Database', 'Employee CSV'))17 18# Predefined queries19predefined_queries = {20    'SQL Database': [21        'Print all students',22        'Count total number of students',23        'List students in Data Science class'24    ],25    'Employee CSV': [26        'Print employees having the department id equal to 100',27        'Count total number of employees',28        'List Top 5 employees according to salary in descending order'29    ]30}31 32st.markdown(f"### Predefined Queries for {data_source}")33 34# Create buttons for predefined queries35for query in predefined_queries[data_source]:36    if st.button(query):37        st.session_state.predefined_query = query38 39st.markdown("### Enter Your Question")40question = st.text_input("Input: ", key="input", value=st.session_state.get('predefined_query', ''))41 42# Submit button43submit = st.button("Submit")44 45if submit:46    # Send request to FastAPI backend47    response = requests.post("http://localhost:8000/query", 48                             json={"question": question, "data_source": data_source})49    if response.status_code == 200:50        data = response.json()51        st.markdown(f"## Generated {'SQL' if data_source == 'SQL Database' else 'Pandas'} Query")52        st.code(data['query'])53        54        st.markdown("## Query Results")55        result = data['result']56        57        if isinstance(result, list) and len(result) > 0:58            if isinstance(result[0], dict):59                # For CSV queries that return a list of dictionaries60                df = pd.DataFrame(result)61                st.dataframe(df)62            elif isinstance(result[0], list):63                # For SQL queries that return a list of lists64                df = pd.DataFrame(result)65                st.dataframe(df)66            else:67                # For single column results68                st.dataframe(pd.DataFrame(result, columns=['Result']))69        elif isinstance(result, dict):70            # For single row results71            st.table(result)72        else:73            # For scalar results or empty results74            st.write(result)75 76        if data_source == 'Employee CSV':77            st.markdown("## Available CSV Columns")78            st.write(data['columns'])79 80        # Update chat history in session state81        st.session_state.chat_history.append(f"👨‍💻({data_source}): {question}")82        st.session_state.chat_history.append(f"🤖: {data['query']}")83    else:84        st.error(f"Error processing your request: {response.text}")85 86    # Clear the predefined query from session state87    st.session_state.pop('predefined_query', None)88 89# Display chat history90st.markdown("## Chat History")91for message in st.session_state.chat_history:92    st.text(message)93 94# Option to clear chat history95if st.button("Clear Chat History"):96    st.session_state.chat_history = []97    st.success("Chat history cleared!")