querybuilder / app.py
Suva's picture
Create app.py
d74e960
raw
history blame
No virus
4.01 kB
# Import required libraries
import os
import openai
import streamlit as st
import pandas as pd
from sqlalchemy import create_engine, text
from PIL import Image
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
temp_db = create_engine('sqlite:///:memory:', echo=True)
# Function to create table definition prompt
def create_table_definition_prompt(col_input):
prompt = '''### sqlite SQL table, with its properties:
#
# Sales({})
#
'''.format(",".join(str(x) for x in col_input))
return prompt
# Function to combine prompts
def combine_prompts(col_input, query_prompt):
definition = create_table_definition_prompt(col_input)
query_init_string = f"### A query to answer: {query_prompt}\nSELECT"
return definition + query_init_string
# Function to execute SQL query
def get_sql(nlp_text, field_input):
print(nlp_text)
print(field_input)
response = openai.Completion.create(
model="text-davinci-003",
prompt=combine_prompts(field_input, nlp_text),
temperature=0,
max_tokens=150,
top_p=1.0,
frequency_penalty=0.0,
presence_penalty=0.0,
stop=["#", ";"]
)
query = response["choices"][0]["text"]
if query.startswith(" "):
query = "SELECT" + query
with temp_db.connect() as conn:
result = conn.execute(text(query))
return result.all()
# Read CSV file into DataFrame
def read_csv_file(uploaded_file):
try:
df = pd.read_csv(uploaded_file, encoding="latin1")
with temp_db.connect() as conn:
df.to_sql(name='Sales', con=conn, if_exists='replace', index=False)
return df
except Exception as e:
st.error(f"Error occurred while reading the CSV file: {str(e)}")
return None
# Configure Streamlit page layout and title
st.set_page_config(
page_title="Query Builder",
layout="wide",
page_icon=":bar_chart:"
)
# Set OpenAI API key
api_key = os.getenv('OPENAI_API_KEY')
#api_key = os.environ.get('OPENAI_API_KEY')
openai.api_key = api_key
# Add CSS styles
st.markdown(
"""
<style>
.front-page {
background-color: #f5f5f5;
padding: 20px;
border-radius: 10px;
}
</style>
""",
unsafe_allow_html=True
)
hide_streamlit_style = """
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
</style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)
col1, col2 = st.columns([1, 2])
with col1:
pass
with col2:
image = Image.open('D:/vscode/qb/data/comsense-Logo2.png')
resized_image = image.resize((1200, 400)) # Adjust the size as per your requirement
st.image(resized_image, width=180)
# Front page content
st.title(":bar_chart: Query Builder")
#st.markdown('<div class="front-page">', unsafe_allow_html=True)
st.title("Please write your *Query* Here :arrow_down_small:")
input_value = st.text_input("", " ")
st.markdown('</div>', unsafe_allow_html=True)
# Sidebar - CSV file upload
uploaded_file = st.sidebar.file_uploader("Upload CSV File")
# Process uploaded CSV file
if uploaded_file is not None:
df = read_csv_file(uploaded_file)
if df is not None:
st.sidebar.dataframe(df)
# Execute SQL query and display result
if st.button("Run Query"):
if uploaded_file is not None and df is not None:
try:
field_input =df.columns
final_result = get_sql(input_value, field_input)
output_column, _ = st.columns(2)
with output_column:
st.title("Output")
if isinstance(final_result, list):
final_result = pd.DataFrame(final_result)
show_data = st.dataframe(final_result.style.set_properties(**{'font-size': '12px', 'text-align': 'center'}))
except Exception as e:
st.error("Please try again with simple sentence.")
else:
st.warning("Please upload a valid CSV file first.")