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