app.py

from dotenv import load_dotenv
import os
from sentence_transformers import SentenceTransformer
import gradio as gr
from sklearn.metrics.pairwise import cosine_similarity
from groq import Groq


load_dotenv()

api = os.getenv("groq_api_key")

def create_metadata_embeddings():
  student = """
  
  Table: student
  Columns:
  - student_id: an integer representing the unique ID of a student.
  - first_name: a string containing the first name of the student.
  - last_name: a string containing the last name of the student.
  - date_of_birth: a date representing the student's birthdate.
  - email: a string for the student's email address.
  - phone_number: a string for the student's contact number.
  - major: a string representing the student's major field of study.
  - year_of_enrollment: an integer for the year the student enrolled.
  - course_id: an integer representing the ID of the course the student is taking (foreign key referencing course_info.course_id).
  """
  professor = """

  Table: professor
  Columns:
  - professor_id: an integer representing the unique ID of a professor.
  - first_name: a string containing the first name of the professor.
  - last_name: a string containing the last name of the professor.
  - email: a string for the professor's email address.
  - department: a string for the department the professor belongs to.
  - position: a string representing the professor's job title.
  - salary: a float representing the professor's salary.
  - date_of_joining: a date for when the professor joined the college.
  """
  
  course = """

  Table: course_info
  Columns:
  - course_id: an integer representing the unique ID of the course.
  - course_name: a string containing the course's name.
  - course_code: a string for the course's unique code.
  - instructor_id: an integer representing the ID of the professor teaching the course (foreign key referencing professor.professor_id).
  - department: a string for the department offering the course.
  - credits: an integer representing the course credits.
  - semester: a string for the semester when the course is offered.
  """



  metadata_list = [student, professor, course]

  model = SentenceTransformer('all-MiniLM-L6-v2')

  embeddings = model.encode(metadata_list)

  return embeddings,model,student,professor,course


def find_best_fit(embeddings,model,user_query,student,professor,course):
  query_embedding = model.encode([user_query])
  similarities = cosine_similarity(query_embedding, embeddings)
  
  table_metadata=""" """
  threshold=similarities[0][similarities.argmax()]*0.8
  table_metadata+=str(threshold)
  if(similarities[0][0]>threshold):
    table_metadata+=student
  if(similarities[0][1]>threshold):
    table_metadata+=professor
  if(similarities[0][2]>threshold):
    table_metadata+=course

  
  
  return table_metadata


def create_prompt(user_query,table_metadata):
  system_prompt="""
  You are a versatile SQL query generator capable of handling natural language queries that may involve a single table or multiple tables joined together. Your task is to accurately interpret user intent and generate valid SQL queries based on the provided metadata.
  
  Rules:
  Dynamic Table Handling: Support both single-table and multi-table queries, ensuring correctness in joins, filters, and aggregations.
  Metadata Validation: Always validate the query against the provided table names, columns, and data types to ensure it is accurate and relevant.
  User Intent: Precisely understand the user's requirements, such as filtering, sorting, aggregations, grouping, or joins, as expressed in natural language.
  Joins: If the query involves multiple tables, use appropriate join conditions based on foreign keys or shared columns provided in the metadata.
  SQL Syntax: Generate SQL queries in standard SQL syntax, ensuring compatibility with most relational database systems.
  
  Output Format:
  Output only the SQL query. Do not include any explanations, comments, or additional text.
  Ensure the entire query is formatted in a single line for simplicity.
  
  Input Format:
  User Query: The natural language description of the required SQL query.
  Table Metadata: The structure of the relevant tables, including table names, column names, data types, and relationships (if applicable).
  
  Output Format:
  A single-line SQL query that adheres to the rules and matches the user's intent.
  You are ready to generate SQL queries based on the user input and provided table metadata.
  """


  user_prompt=f"""
  User Query: {user_query}
  Table Metadata: {table_metadata}
  """

  return system_prompt,user_prompt


def generate_output(system_prompt,user_prompt):
  client = Groq(api_key=api,)
  chat_completion = client.chat.completions.create(messages=[
    {"role": "system", "content": system_prompt},
     {"role": "user","content": user_prompt,}],model="llama3-8b-8192",)
  res = chat_completion.choices[0].message.content

  select=res[0:6].lower()
  if(select=="select"):
    output=res
  else:
    output="Can't perform the task at the moment."
  

  return output

def response(user_query):
  embeddings,model,student,professor,course=create_metadata_embeddings()

  table_metadata=find_best_fit(embeddings,model,user_query,student,professor,course)

  system_prompt,user_prompt=create_prompt(user_query,table_metadata)

  output=generate_output(system_prompt,user_prompt)

  return output


desc="""
There are three tables in the database:

Student Table:
This table contains information about students, including the student's unique ID, first name, last name, date of birth, email address, phone number, major field of study, year of enrollment, and the ID of the course they are enrolled in. The course_id serves as a foreign key referencing the course_info table, linking each student to a single course.

Professor Table:
This table includes details about professors, such as the professor's unique ID, first name, last name, email address, department, job position, salary, and date of joining.

Course Info Table:
This table stores details about courses, including the course's unique ID, name, course code, instructor ID, department offering the course, number of credits, and the semester in which the course is offered. The instructor_id is a foreign key referencing the professor table, associating each course with the professor who teaches it.

"""


demo = gr.Interface(
    fn=response,
    inputs=gr.Textbox(label="Please provide the natural language query"),
    outputs=gr.Textbox(label="SQL Query"),
    title="SQL Query generator",
    description=desc
)

demo.launch()