Arnav-J/Schema-SQL-Query

1from types import ClassMethodDescriptorType2from click import prompt3from dotenv import load_dotenv4from langchain.llms import OpenAI5from langchain import PromptTemplate6from langchain.chains import LLMChain7import streamlit as st8from streamlit_extras.let_it_rain import rain9import os10import time11import random 12 13#setting up the prerequites14 15# load_dotenv()16# os.environ["OPENAI_API_KEY"] = os.getenv('openai_key')17llm = OpenAI(temperature = 0.8)18LANGCHAIN_TRACING_V2="true"19# os.environ["LANGCHAIN_API_KEY"] = os.getenv('LANGCHAIN_API_KEY')20# os.environ["LANGCHAIN_TRACING_V2"] = "true"21#llm2 = OpenAI(model_name = "gpt-4-turbo-preview", temperature = 0.7)22 23st.set_page_config(page_title="Schem-SQL-Query",24    page_icon="🧊",25    layout="wide",26    initial_sidebar_state="collapsed")27 28rain(29    emoji="❄️",30    font_size=20,31    falling_speed=10,32    animation_length="infinite",33)34 35with open( "style.css" ) as css:36    st.markdown( f'<style>{css.read()}</style>' , unsafe_allow_html= True)37 38st.markdown("<h1 style='text-align: center; color: white; text-decoration: underline; font-family: 'Times New Roman';'>SQL Generator</h1>", unsafe_allow_html=True)39col1 , col2 = st.columns([0.4,0.6], gap="large")40 41col1.header("Database Schema Input")42schema_input = col1.text_area("Submit Database Schema", height=450)43schema_submit_button = col1.button("Submit Schema", key="schema_submit")44 45col2.header("Enter your question")46user_prompt = col2.text_area("User's Prompt", height=150)47prompt_submit_button = col2.button("Submit Prompt", key="prompt_submit")48 49col2.markdown("---")50 51#query prompt52q_template = """You are an expert in writing SQL queries thus, based on the table schema below, 53write a SQL query that would answer the user's question:54{schema} \n Question: {user_prompt} \nSQL Query:"""55 56## Prompt Templates57first_input_prompt = PromptTemplate(58    input_variables = ['schema','user_prompt'],59    template = q_template60)61 62#query chain63chain1  = LLMChain(llm = llm, prompt = first_input_prompt, verbose=True, output_key='query')64########################65 66#scenario prompt67s_template = """Based on the following schema, generate 9 sql questions that 68are relevant to the schema. The questions should be in increasing order of difficulty and complexity.69Include atleast one problem which involves "join" operation. 70Just describe the problem statement in one line. Schema: {schema}  """71 72third_input_prompt = PromptTemplate(73    input_variables = ['schema'],74    template = s_template75)76 77#scenario chain78chain3 = LLMChain(llm = llm, prompt = third_input_prompt, verbose=True, output_key='scenario')79 80if schema_submit_button:81    with col1:82        with st.spinner('Wait for it...'):83            time.sleep(random.random() * 3)84        col1.success("Schema Submitted Successfully!!")85    86    87    expander = col1.expander("See Scenarios")88    scenarios = chain3({"schema": schema_input}, return_only_outputs=True)89    expander.write(scenarios["scenario"])90      91 92############################93    94#algorithm prompt95a_template = """You are an expert in explaining SQL queries thus96Explain the following SQL Query in an Algorithmic step-by-step fashion. Write "Start" and "End" to denote the begining 97and ending of the Algorithm. Write at least 4 lines and a maximum of 12 lines to explain.98Explain each step briefly, in a single line. Refer the below example.99Example : Start1001. Start by selecting distinct customer details and product application details.1012. Join the customers table with the life events table on customer id.1023. Join the resulting set with the products table on customer_id.1034. Filter for customers who are household heads.1045. Filter for customers located in California.1056. Filter for life events that include 'relocation'.1067. Filter for life events that occurred in the last 6 months.1078. Filter for product applications that include 'insurance'.1089. Filter for product applications that occurred in the last 6 months.109\n110End111Query: {query}"""112 113second_input_prompt = PromptTemplate(114    input_variables = ['query'],115    template = a_template116)117 118#algorithm chain119chain2  = LLMChain(llm = llm, prompt = second_input_prompt, verbose=True, output_key='algorithm')120 121if prompt_submit_button:122    with st.spinner('Generating response...'):123            qu = chain1({"schema": schema_input, "user_prompt":user_prompt}, return_only_outputs=True)124            #print(type(qu))125            col2.subheader("Query")126            col2.code(qu["query"], language="sql")127            #st.write(qu["query"])128            col2.markdown("---")129            col2.subheader("Algorithm")130            algo = chain2({"query": qu}, return_only_outputs=True)131            col2.write(algo["algorithm"])