AI-SQL-query-generator/sql.py at main · janampatel/AI-SQL-query-generator · GitHub

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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import re
import os
import json
import openai
import pathlib
import streamlit as st
from openai import OpenAI
from streamlit_lottie import st_lottie
import sqlite3

client = OpenAI(base_url="http://localhost:1234/v1", api_key="lm-studio")


st.set_page_config(
    page_title="SQL Query Generator",
    page_icon="✨",
    layout= "wide",
    initial_sidebar_state="expanded",
    menu_items={
    'Get Help': 'https://github.com/mann1105',
    'Report a bug': "mailto:[email protected]",
    'About': "## A minimalistic application to generate SQL queries using OpenAI APIs built with Python and Streamlit"
    } )

@st.cache_data()
def lottie_local(filepath: str):
    with open(filepath, "r") as f:
        return json.load(f)

@st.cache_data()
def hide_footer():
    hide_st_style = """
            <style>
            footer {visibility: hidden;}
            </style>
            """
    st.markdown(hide_st_style, unsafe_allow_html=True)

st.title("Natural Language <2> SQL Query 🚀")
hide_footer()

col1, col2 = st.columns(2)
with col1:
    anim = lottie_local(r"./assets/animation.json")
    st_lottie(anim,
            speed=1,
            reverse=False,
            loop=True,
            height = 700,
            width = 700,
            quality="high",
            key=None)

with col2:
    st.markdown("-----------------------------")
    english_input = st.text_area("Please enter the desired question and see the magic 💫", height = 250)

    prompt = f"Translate this natural language query into syntactically correct SQL:\n\n{english_input}\n\nSQL Query:"
    st.markdown("-----------------------------")
    ch = st.checkbox("Table Schema")
    if ch:
        schema = st.text_area("Enter Table Schema 📝")
        # prompt = f"Translate this natural language query into syntactically correct SQL:\n\n{query}\n\nUse this table schema:\n\n{schema}\n\n{prompt}"
        prompt = f"""
        ### Task

        Generate a SQL query to answer the following question: {english_input}

        ### PostgreSQL Database Schema

        The query will run on a database with the following schema:

        {schema}

        ### Answer

        Here is the SQL query that answers the question: {english_input}

        ```sql

        """
        st.markdown("-----------------------------")
    # if st.button("Generate SQL Query ✨", use_container_width=True):
    #     with st.spinner("Working.. 💫"):
    #         try:
    #             completion = client.chat.completions.create(
    #                 model = "chatdb/natural-sql-7b-GGUF",
    #                 messages = [
    #                     {"role": "user", "content":prompt}
    #                 ]
    #             )

    #             response = completion.choices[0].text
    #             st.balloons()
    #             st.markdown("### Output:")
    #             st.success(f"{response}")
    #         except Exception as e:
    #             st.error(f"Error: {e}")
    if st.button("Generate SQL Query ✨", use_container_width=True):
        with st.spinner("Working.." ):
            try:
                 completion = client.chat.completions.create(
                # ...
            )
                 response = completion.choices[0].text
                 st.balloons()
                 st.markdown("## Output:")

            # Connect to SQLite database
                 conn = sqlite3.connect('your_database.db')
                 cursor = conn.cursor()

            # Execute the query
                 cursor.execute(response)

            # Fetch and display results (if applicable)
                 results = cursor.fetchall()
                 st.success(f"{response}\n")
                 if results:
                   st.dataframe(results,  width=1000, height=400)

            # Close connection
                 conn.close()

            except Exception as e:
             st.error(f"Error: {e}")