Upload 9 files

app.py

@@ -1,23 +1,23 @@
-import gradio as gr
-from db.database import Database
-from nlp.query_processor import QueryProcessor
-
-# Initialize Database and Query Processor
-db = Database(db_name="chat_assistant.db")
-query_processor = QueryProcessor(db)
-
-def respond(message, history):
-    """Processes user queries, fetches results from the database, and returns responses."""
-    response = query_processor.process_query(message)
-    return response
-
-# Gradio Chat UI
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[],
-    title="SQL Chat Assistant",
-    description="Ask any database-related question, and I will generate an SQL query and fetch the relevant data.",
-)
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+from db.database import Database
+from nlp.query_processor import QueryProcessor
+
+# Initialize Database and Query Processor
+db = Database(db_name="chat_assistant.db")
+query_processor = QueryProcessor(db)
+
+def respond(message, history):
+    """Processes user queries, fetches results from the database, and returns responses."""
+    response = query_processor.process_query(message)
+    return response
+
+# Gradio Chat UI
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[],
+    title="SQL Chat Assistant",
+    description="Ask any database-related question, and I will generate an SQL query and fetch the relevant data.",
+)
+
+if __name__ == "__main__":
+    demo.launch()

db/database.py

@@ -0,0 +1,31 @@
+import sqlite3
+
+class Database:
+    def __init__(self, db_name):
+        self.db_name = db_name
+        self.conn = None
+
+    def connect(self):
+        try:
+            self.conn = sqlite3.connect(self.db_name)
+            return self.conn
+        except sqlite3.Error as e:
+            print(f"Database connection error: {e}")
+            return None
+
+    def close(self):
+        if self.conn:
+            self.conn.close()
+
+    def execute_query(self, query, params=None):
+        try:
+            cursor = self.conn.cursor()
+            if params:
+                cursor.execute(query, params)
+            else:
+                cursor.execute(query)
+            self.conn.commit()
+            return cursor
+        except sqlite3.Error as e:
+            print(f"Query execution error: {e}")
+            return None

db/init_db.py

@@ -0,0 +1,77 @@
+import sqlite3
+
+def initialize_database(db_name='chat_assistant.db'):
+    conn = None
+    try:
+        conn = sqlite3.connect(db_name)
+        cursor = conn.cursor()
+
+        # Create Employees table
+        cursor.execute('''CREATE TABLE IF NOT EXISTS employees (
+            ID INTEGER PRIMARY KEY,
+            Name TEXT,
+            Department TEXT,
+            Salary REAL,
+            Hire_Date TEXT
+        )''')
+
+        # Create Departments table
+        cursor.execute('''CREATE TABLE IF NOT EXISTS departments (
+            ID INTEGER PRIMARY KEY,
+            Name TEXT,
+            Manager TEXT
+        )''')
+
+        cursor.execute('''CREATE TABLE IF NOT EXISTS table_metadata (
+            table_name TEXT PRIMARY KEY,
+            description TEXT
+        )''')
+
+        cursor.execute('''CREATE TABLE IF NOT EXISTS column_metadata (
+            table_name TEXT,
+            column_name TEXT,
+            data_type TEXT,
+            description TEXT,
+            PRIMARY KEY (table_name, column_name),
+            FOREIGN KEY (table_name) references table_metadata (table_name)
+            )''')
+        # Insert sample data into Employees table
+        employees_data = [
+            (1, 'Alice', 'Sales', 50000, '2021-01-15'),
+            (2, 'Bob', 'Engineering', 70000, '2020-06-10'),
+            (3, 'Charlie', 'Marketing', 60000, '2022-03-20')
+        ]
+        cursor.executemany('INSERT INTO Employees VALUES (?, ?, ?, ?, ?)', employees_data)
+
+        # Insert sample data into Departments table
+        departments_data = [
+            (1, 'Sales', 'Alice'),
+            (2, 'Engineering', 'Bob'),
+            (3, 'Marketing', 'Charlie')
+        ]
+        cursor.executemany('INSERT INTO Departments VALUES (?, ?, ?)', departments_data)
+        cursor.execute("INSERT INTO table_metadata (table_name, description) VALUES ('Employees','Details of Employees in Department.')")
+        cursor.execute("INSERT INTO table_metadata (table_name, description) VALUES ('Department','Details of Manager of the Department.')")
+
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Employees','id','INTEGER','Identification number of the Employee.')")
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Employees','name','TEXT','Name of the Employee.')")
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Employees','department','TEXT','Department of the Employee.')")
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Employees','Salary','INTEGER','Salary of the Employee.')")
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Employees','Hire_Date','DATE','Date in which the Employee was hired.')")
+
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Departments','id','INTEGER','Identification number of the Employee.')")
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Departments','name','TEXT','Name of the Department.')")
+        cursor.execute("INSERT INTO column_metadata (table_name, column_name, data_type,description) VALUES ('Departments','manager','TEXT','Manager of the Department.')")
+        conn.commit()
+        
+        print("Database initialized successfully.")
+
+    except sqlite3.Error as e:
+        print(f"Database initialization error: {e}")
+
+    finally:
+        if conn:
+            conn.close()
+
+if __name__ == "__main__":
+    initialize_database()

nlp/query_processor.py

@@ -0,0 +1,60 @@
+import re
+import sqlite3  # Added import for sqlite3
+from datetime import datetime
+from db.database import Database
+from nlp.sql_model import SQLModel
+
+class QueryProcessor:
+    def __init__(self, db):
+        self.db = db
+        self.sql_model = SQLModel()
+
+    def convert_date(self, date_str):
+        formats = [
+            '%d/%m/%Y',
+            '%d %B %Y',
+            '%B %Y',
+            '%b %Y',
+            '%Y-%m-%d'
+        ]
+        for fmt in formats:
+            try:
+                return datetime.strptime(date_str, fmt).date()
+            except ValueError:
+                continue
+        raise ValueError('Invalid date format')
+
+    def extract_and_convert_date(self, user_query):
+        date_pattern = r'\b(\b[0-9]{1,2}[/-]?[0-9]{1,2}[/-]?[0-9]{2,4}\b|\b(?:\b(?:January|February|March|April|May|June|July|August|September|October|November|December|Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\b[- ]?[0-9]{1,2}[- ]?[0-9]{2,4}\b|\b[0-9]{2,4}[- ]?(?:January|February|March|April|May|June|July|August|September|October|November|December|Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\b)\b)'
+        match = re.search(date_pattern, user_query, re.I)
+        if match:
+            date_str = match.group(0)
+            try:
+                return self.convert_date(date_str)
+            except ValueError:
+                return None
+        return None
+
+
+    def process_query(self, user_query):
+        conn = self.db.connect()
+        if not conn:
+            return "Failed to connect to the database."
+
+        response = ''
+        try:
+            # Generate SQL query using the LLM model
+            sql_query = self.sql_model.generate_sql(user_query)
+            # Execute generated SQL query against the database
+            cursor = self.db.execute_query(sql_query)
+            if cursor:
+                results = cursor.fetchall()
+                if results:
+                    response = f'Results: {results}'
+                else:
+                    response = 'No results found.'
+            else:
+                response = 'Error executing query.'
+        finally:
+            self.db.close()
+        return response

nlp/sql_model.py

@@ -0,0 +1,111 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+
+class SQLModel:
+    def __init__(self, model_name="google/flan-t5-base"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        # self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
+        # self.model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it",)
+
+    def generate_sql(self, natural_language_query):
+        input_text = f"""You are a highly skilled SQL translator. Your task is to convert natural language descriptions of data queries into correct and optimized SQL statements.
+        Here is the schema information for our database :
+
+        Table: Employees
+        - id (INT)
+        - NAME (VARCHAR)
+        - Department (VARCHAR)
+        - Salary (INT)
+        - Hire_Date (DATE)
+
+        Table: Departments
+        - ID (INT)
+        - Name (VARCHAR)
+        - Manager (VARCHAR)
+
+        Here are a few examples:
+
+        1. **Input**: "Show me all employees in the Sales department."
+        **Output**:
+
+                SELECT *
+                FROM Employees
+                WHERE Department = 'Sales';
+
+        2. **Input**: "Who is the manager of the Engineering department?"
+        **Output**:
+
+                SELECT Manager
+                FROM Departments
+                WHERE Name = 'Engineering';
+
+
+        3. **Input**: "List all employees hired after 2021-01-01."
+        **Output**:
+
+                SELECT *
+                FROM Employees
+                WHERE Hire_Date > '2021-01-01';
+
+
+        4. **Input**: "What is the total salary expense for the Marketing department?"
+        **Output**:
+
+                SELECT SUM(Salary)
+                FROM Employees
+                WHERE Department = 'Marketing';
+
+
+        5. **Input**: "Find the average salary of employees in each department."
+        **Output**:
+
+                SELECT Department, AVG(Salary) AS average_salary
+                FROM Employees
+                GROUP BY Department;
+
+        Please do not return additional text besides query.
+
+        Please only answer queries which makes sense for the given schema. Else just return - "No information found"
+
+        Now, translate the following natural language query into an syntactically correct SQL query:
+        **Input**: {natural_language_query}
+        **Output**:
+
+"""
+#         input_text = f"""
+# Translate the following natural language query into a syntactically correct SQL query using the provided database schema. Output only the SQL query with no additional text or explanation.
+
+# Database Schema:
+
+# Table: Employees
+# - id (INT)
+# - NAME (VARCHAR)
+# - Department (VARCHAR)
+# - Salary (INT)
+# - Hire_Date (DATE)
+
+# Table: Departments
+# - ID (INT)
+# - Name (VARCHAR)
+# - Manager (VARCHAR)
+
+# Examples:
+# 1. Natural Language Query: "List all employees who were hired after '2020-01-01'."
+#    Output: SELECT * FROM Employees WHERE Hire_Date > '2020-01-01';
+
+# 2. Natural Language Query: "Retrieve the names and salaries of employees in the 'Sales' department."
+#    Output: SELECT NAME, Salary FROM Employees WHERE Department = 'Sales';
+
+# Now, translate the following query:
+# {natural_language_query}
+# """
+        # input_text = f"translate English to SQL: {natural_language_query}"
+        # inputs = self.tokenizer(input_text, return_tensors="pt").input_ids
+        # outputs = self.model.generate(inputs, max_new_tokens=100, do_sample=False)
+        # sql_query = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        inputs = self.tokenizer(input_text, return_tensors="pt")
+        outputs = self.model.generate(**inputs)
+        sql_query = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print(sql_query)
+        return sql_query