text2sql_ui.py

"""
UI for text2sql app
"""
import os
import pandas as pd
import requests
import streamlit as st

# Streamlit app
st.set_page_config(layout="wide")


def main():
    st.title("Mutual Fund Text2SQL App")

    # Get user prompt from Streamlit UI
    prompt = st.text_input("Enter your question here:")

    if st.button("Submit"):

        API_URL = f"{os.environ['SERVER_URL']}/api/get-mf-data/?query={prompt}"
        response = requests.get(API_URL)

        if response.status_code != 200:
            st.error("Error fetching data from the server.")
            st.stop()

        df = pd.DataFrame(response.json()["data"])
        st.write("Query:", response.json()["query"])
        # st.markdown(
        #     "<h1 style='text-align: center;'>Mutual Fund Data Analysis Tool</h1>",
        #     unsafe_allow_html=True,
        # )

        # Display the DataFrame without scrolling and use the full page width
        st.dataframe(df, width=10000, height=1000)


if __name__ == "__main__":
    main()