Upload 9 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Chinook_Sqlite.sqlite filter=lfs diff=lfs merge=lfs -text

Chinook_Sqlite.sqlite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdf635be69850bd3be09c9a2dbeef7ddfb80036bd3ef3381383cd03b61e4a61a
+size 1067008

app.py

@@ -1,5 +1,63 @@
 import streamlit as st
+from lib import Text2SQLRAG
+from utils import execute_query_and_return_df
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
 
+st.set_page_config(page_title="Text2SQLRAG")
+st.title("Text2SQLRAG")
+
+# Create an instance of Text2SQLRAG
+text2sql = Text2SQLRAG()
+
+# Initialize session state for storing chat messages
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+    
+    # Display conversation history from session state
+for message in st.session_state.messages:
+    role = message.get("role", "assistant")
+    with st.chat_message(role):
+        if "output" in message:
+            with st.expander("Reasoning", expanded=True):
+                st.markdown(message["reasoning"])
+        if "sql_query" in message and message["sql_query"]:
+            with st.expander("SQL Query", expanded=True):
+                st.code(message["sql_query"])
+
+# Get user input
+input_text = st.chat_input("Chat with your bot here...")
+
+if input_text:
+    # Display user input
+    with st.chat_message("user"):
+        st.markdown(input_text)
+
+    # Add user input to chat history
+    st.session_state.chat_history.append({"role": "user", "text": input_text})
+
+    # Get chatbot response
+    response = text2sql.run(input_text)
+    sql_query = response.query
+    reasoning = response.reasoning
+    df = execute_query_and_return_df(sql_query)
+
+    if sql_query:
+        with st.expander("SQL Query", expanded=True):
+            st.code(sql_query)
+
+        with st.expander("Reasoning", expanded=True):
+            st.write(reasoning)
+
+    if df is not None:
+        st.dataframe(df)
+    else:
+        st.error("Error executing query")
+    
+    # Append assistant response to session state
+    st.session_state.messages.append(
+        {
+            "role": "assistant",
+            "reasoning": reasoning,
+            "sql_query": sql_query,
+        }
+    )

constant.py

@@ -0,0 +1,10 @@
+import os
+from dotenv import load_dotenv, find_dotenv
+load_dotenv(find_dotenv())
+
+
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+EMBEDDING_MODEL = "intfloat/multilingual-e5-small"
+GEMINI_MODEL = "gemini-2.0-flash"
+PATH_SCHEMA = "/data/music-database-schema.json"
+PATH_DB = "/data/Cinhook_Sqlite.sqlite"

lib.py

@@ -0,0 +1,102 @@
+import logging
+from pydantic import BaseModel, Field
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.runnables import RunnablePassthrough
+from langchain_chroma import Chroma
+from langchain_huggingface.embeddings import HuggingFaceEmbeddings
+from langchain_community.document_loaders import JSONLoader  
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+from prompt import prompt
+from utils import execute_query_and_return_df
+from constant import (
+    GEMINI_MODEL, 
+    GOOGLE_API_KEY, 
+    PATH_SCHEMA, 
+    PATH_DB,
+    EMBEDDING_MODEL
+)
+
+
+
+class SQLOutput(BaseModel):
+    query: str = Field(description="The SQL query to run.")
+    reasoning: str = Field(description="Reasoning to understand the SQL query.")
+
+class Text2SQLRAG:
+    def __init__(self, 
+                path_schema: str = PATH_SCHEMA,
+                path_db: str = PATH_DB,
+                model: str = GEMINI_MODEL, 
+                api_key: str = GOOGLE_API_KEY,
+                embedding_model: str = EMBEDDING_MODEL
+            ):
+        """
+        A class for generating SQL queries based on natural language text.
+        """
+        self.logger = logging.getLogger(__name__)
+        self.logger.info('Initializing Text2SQLRAG')
+
+        model_kwargs = {  
+            "max_tokens": 512,  
+            "temperature": 0.2,  
+            "top_k": 250,  
+            "top_p": 1,  
+            "stop_sequences": ["\n\nHuman:"] 
+        }
+        self.schema = path_schema
+        self.db = path_db
+        self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
+        
+        self.model = ChatGoogleGenerativeAI(
+            model=model,
+            api_key=api_key,
+            model_kwargs=model_kwargs
+        )
+        self.llm = self.model.with_structured_output(SQLOutput)
+        self.retriever = self._indexing_vectore()
+
+    def _indexing_vectore(self):
+        """
+        Indexes the database schema using a vector store for efficient retrieval.
+
+        This method loads the schema from a JSON file, splits it into chunks,
+        embeds the chunks using a specified embedding model, and stores them in
+        a vector store. It returns a retriever configured to search for the top
+        k relevant documents.
+
+        Returns:
+            retriever: An object capable of retrieving the most relevant schema
+            chunks based on the given search parameters.
+        """
+
+        self.logger.info('Indexing schema')
+        db_schema_loader = JSONLoader(
+            file_path=self.schema,  
+            jq_schema='.',  
+            text_content=False
+        )  
+        text_splitter = RecursiveCharacterTextSplitter(  
+            separators=["separator"], 
+            chunk_size=10000,  
+            chunk_overlap=100  
+        )
+
+        docs = text_splitter.split_documents(db_schema_loader.load())  
+
+        vectorstore = Chroma.from_documents(documents=docs, 
+                                        embedding=self.embeddings)
+
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
+        self.logger.info('Finished indexing schema')
+        return retriever
+    
+    def run(self, question: str):
+        self.logger.info(f'Running Text2SQLRAG for question: {question}')
+        rag_chain = (
+            {"context": self.retriever, "question": RunnablePassthrough()}
+            | prompt
+            | self.llm
+        )
+        response = rag_chain.invoke(question)
+        return response

music-database-schema.json

@@ -0,0 +1,757 @@
+{
+    "tables": [
+        {
+            "separator": "table_1",
+            "name": "Album",
+            "schema": "CREATE TABLE Album (AlbumId integer PRIMARY KEY, Title character varying(160), ArtistId integer);",
+            "description": "This table stores information about music albums.",
+            "columns": [
+                {
+                    "name": "AlbumId",
+                    "description": "unique identifier for albums.",
+                    "synonyms": [
+                        "album id"
+                    ]
+                },
+                {
+                    "name": "Title",
+                    "description": "title of the album",
+                    "synonyms": [
+                        "album title",
+                        "album name"
+                    ]
+                },
+                {
+                    "name": "ArtistId",
+                    "description": "Id of the artist associated with the album",
+                    "synonyms": [
+                        "artist id"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT a.Title FROM Album a JOIN Artist ar ON a.ArtistId = ar.ArtistId WHERE ar.Name = 'AC/DC'",
+                    "user_input": "Get all albums by AC/DC"
+                },
+                {
+                    "query": "SELECT COUNT(*) FROM Album WHERE ArtistId = 1",
+                    "user_input": "Count how many albums AC/DC has"
+                },
+                {
+                    "query": "SELECT a.Title, ar.Name FROM Album a JOIN Artist ar ON a.ArtistId = ar.ArtistId ORDER BY a.Title ASC",
+                    "user_input": "List all albums alphabetically with their artists"
+                },
+                {
+                    "query": "SELECT a.Title FROM Album a JOIN Track t ON a.AlbumId = t.AlbumId GROUP BY a.AlbumId ORDER BY COUNT(t.TrackId) DESC LIMIT 1",
+                    "user_input": "Find the album with the most tracks"
+                }
+            ]
+        },
+        {
+            "separator": "table_2",
+            "name": "Artist",
+            "schema": "CREATE TABLE Artist (ArtistId integer PRIMARY KEY, Name character varying(120));",
+            "description": "This table stores information about music artists.",
+            "columns": [
+                {
+                    "name": "ArtistId",
+                    "description": "unique identifier for artists.",
+                    "synonyms": [
+                        "artist id"
+                    ]
+                },
+                {
+                    "name": "Name",
+                    "description": "name of the artist",
+                    "synonyms": [
+                        "artist name"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT Name FROM Artist WHERE ArtistId = 1",
+                    "user_input": "Get artist name with id 1"
+                },
+                {
+                    "query": "SELECT ar.Name, COUNT(a.AlbumId) as AlbumCount FROM Artist ar LEFT JOIN Album a ON ar.ArtistId = a.ArtistId GROUP BY ar.ArtistId ORDER BY AlbumCount DESC",
+                    "user_input": "List artists by number of albums they have"
+                },
+                {
+                    "query": "SELECT ar.Name FROM Artist ar JOIN Album a ON ar.ArtistId = a.ArtistId JOIN Track t ON a.AlbumId = t.AlbumId GROUP BY ar.ArtistId ORDER BY SUM(t.Milliseconds) DESC LIMIT 5",
+                    "user_input": "Find top 5 artists with the longest total music duration"
+                },
+                {
+                    "query": "SELECT ar.Name FROM Artist ar JOIN Album a ON ar.ArtistId = a.ArtistId JOIN Track t ON a.AlbumId = t.AlbumId JOIN InvoiceLine il ON t.TrackId = il.TrackId GROUP BY ar.ArtistId ORDER BY SUM(il.UnitPrice * il.Quantity) DESC LIMIT 3",
+                    "user_input": "Find the top 3 best-selling artists by revenue"
+                }
+            ]
+        },
+        {
+            "separator": "table_3",
+            "name": "Customer",
+            "schema": "CREATE TABLE Customer (CustomerId integer PRIMARY KEY, FirstName character varying(40), LastName character varying(20), Company character varying(80), Address character varying(70), City character varying(40), State character varying(40), Country character varying(40), PostalCode character varying(10), Phone character varying(24), Fax character varying(24), Email character varying(60), SupportRepId integer);",
+            "description": "This table stores information about customers who purchase music.",
+            "columns": [
+                {
+                    "name": "CustomerId",
+                    "description": "unique identifier for customers",
+                    "synonyms": [
+                        "customer id"
+                    ]
+                },
+                {
+                    "name": "FirstName",
+                    "description": "first name of the customer",
+                    "synonyms": [
+                        "first name"
+                    ]
+                },
+                {
+                    "name": "LastName",
+                    "description": "last name of the customer",
+                    "synonyms": [
+                        "last name",
+                        "surname"
+                    ]
+                },
+                {
+                    "name": "Company",
+                    "description": "company name of the customer if applicable",
+                    "synonyms": [
+                        "company name",
+                        "organization"
+                    ]
+                },
+                {
+                    "name": "Address",
+                    "description": "street address of the customer",
+                    "synonyms": [
+                        "street address"
+                    ]
+                },
+                {
+                    "name": "City",
+                    "description": "city where the customer lives",
+                    "synonyms": [
+                        "customer city"
+                    ]
+                },
+                {
+                    "name": "State",
+                    "description": "state or province where customer lives",
+                    "synonyms": [
+                        "province"
+                    ]
+                },
+                {
+                    "name": "Country",
+                    "description": "country where customer lives",
+                    "synonyms": [
+                        "nation"
+                    ]
+                },
+                {
+                    "name": "PostalCode",
+                    "description": "postal or zip code of customer's address",
+                    "synonyms": [
+                        "zip code"
+                    ]
+                },
+                {
+                    "name": "Phone",
+                    "description": "phone number of the customer",
+                    "synonyms": [
+                        "telephone"
+                    ]
+                },
+                {
+                    "name": "Fax",
+                    "description": "fax number of the customer if available",
+                    "synonyms": [
+                        "facsimile"
+                    ]
+                },
+                {
+                    "name": "Email",
+                    "description": "email address of the customer",
+                    "synonyms": [
+                        "email address"
+                    ]
+                },
+                {
+                    "name": "SupportRepId",
+                    "description": "Id of employee assigned as the customer's support representative",
+                    "synonyms": [
+                        "support rep id",
+                        "employee id"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT FirstName, LastName, Email FROM Customer WHERE Country = 'Brazil'",
+                    "user_input": "Get all Brazilian customers with their emails"
+                },
+                {
+                    "query": "SELECT Country, COUNT(*) as CustomerCount FROM Customer GROUP BY Country ORDER BY CustomerCount DESC",
+                    "user_input": "List countries by number of customers"
+                },
+                {
+                    "query": "SELECT c.FirstName, c.LastName, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.CustomerId ORDER BY TotalSpent DESC LIMIT 5",
+                    "user_input": "Find top 5 customers by total spending"
+                },
+                {
+                    "query": "SELECT c.Email FROM Customer c LEFT JOIN Invoice i ON c.CustomerId = i.CustomerId WHERE i.InvoiceId IS NULL",
+                    "user_input": "Find customers who haven't made any purchases"
+                }
+            ]
+        },
+        {
+            "separator": "table_4",
+            "name": "Employee",
+            "schema": "CREATE TABLE Employee (EmployeeId integer PRIMARY KEY, LastName character varying(20), FirstName character varying(20), Title character varying(30), ReportsTo integer, BirthDate timestamp without time zone, HireDate timestamp without time zone, Address character varying(70), City character varying(40), State character varying(40), Country character varying(40), PostalCode character varying(10), Phone character varying(24), Fax character varying(24), Email character varying(60));",
+            "description": "This table stores information about employees who work for the music store.",
+            "columns": [
+                {
+                    "name": "EmployeeId",
+                    "description": "unique identifier for employees",
+                    "synonyms": [
+                        "employee id"
+                    ]
+                },
+                {
+                    "name": "LastName",
+                    "description": "last name of the employee",
+                    "synonyms": [
+                        "surname"
+                    ]
+                },
+                {
+                    "name": "FirstName",
+                    "description": "first name of the employee",
+                    "synonyms": [
+                        "given name"
+                    ]
+                },
+                {
+                    "name": "Title",
+                    "description": "job title of the employee",
+                    "synonyms": [
+                        "position",
+                        "job title"
+                    ]
+                },
+                {
+                    "name": "ReportsTo",
+                    "description": "Id of the employee's manager",
+                    "synonyms": [
+                        "manager id",
+                        "supervisor id"
+                    ]
+                },
+                {
+                    "name": "BirthDate",
+                    "description": "date of birth of the employee",
+                    "synonyms": [
+                        "date of birth",
+                        "DOB"
+                    ]
+                },
+                {
+                    "name": "HireDate",
+                    "description": "date when the employee was hired",
+                    "synonyms": [
+                        "start date",
+                        "employment date"
+                    ]
+                },
+                {
+                    "name": "Address",
+                    "description": "street address of the employee",
+                    "synonyms": [
+                        "street address"
+                    ]
+                },
+                {
+                    "name": "City",
+                    "description": "city where the employee lives",
+                    "synonyms": [
+                        "employee city"
+                    ]
+                },
+                {
+                    "name": "State",
+                    "description": "state or province where employee lives",
+                    "synonyms": [
+                        "province"
+                    ]
+                },
+                {
+                    "name": "Country",
+                    "description": "country where employee lives",
+                    "synonyms": [
+                        "nation"
+                    ]
+                },
+                {
+                    "name": "PostalCode",
+                    "description": "postal or zip code of employee's address",
+                    "synonyms": [
+                        "zip code"
+                    ]
+                },
+                {
+                    "name": "Phone",
+                    "description": "phone number of the employee",
+                    "synonyms": [
+                        "telephone"
+                    ]
+                },
+                {
+                    "name": "Fax",
+                    "description": "fax number of the employee if available",
+                    "synonyms": [
+                        "facsimile"
+                    ]
+                },
+                {
+                    "name": "Email",
+                    "description": "email address of the employee",
+                    "synonyms": [
+                        "email address"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT FirstName, LastName FROM Employee WHERE Title = 'Sales Manager'",
+                    "user_input": "Get all sales managers"
+                },
+                {
+                    "query": "SELECT e1.FirstName, e1.LastName, e2.FirstName as ManagerFirstName, e2.LastName as ManagerLastName FROM Employee e1 LEFT JOIN Employee e2 ON e1.ReportsTo = e2.EmployeeId",
+                    "user_input": "List all employees with their managers"
+                },
+                {
+                    "query": "SELECT COUNT(*) FROM Employee WHERE STRFTIME('%Y', HireDate) = '2002'",
+                    "user_input": "Count employees hired in 2002"
+                },
+                {
+                    "query": "SELECT e.FirstName, e.LastName, SUM(i.Total) as TotalSales FROM Employee e JOIN Customer c ON e.EmployeeId = c.SupportRepId JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY e.EmployeeId ORDER BY TotalSales DESC",
+                    "user_input": "Rank employees by total sales from their customers"
+                }
+            ]
+        },
+        {
+            "separator": "table_5",
+            "name": "Genre",
+            "schema": "CREATE TABLE Genre (GenreId integer PRIMARY KEY, Name character varying(120));",
+            "description": "This table stores music genres.",
+            "columns": [
+                {
+                    "name": "GenreId",
+                    "description": "unique identifier for genres",
+                    "synonyms": [
+                        "genre id"
+                    ]
+                },
+                {
+                    "name": "Name",
+                    "description": "name of the genre",
+                    "synonyms": [
+                        "genre name"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT Name FROM Genre WHERE GenreId = 1",
+                    "user_input": "Get the genre name with id 1"
+                },
+                {
+                    "query": "SELECT g.Name, COUNT(t.TrackId) as TrackCount FROM Genre g JOIN Track t ON g.GenreId = t.GenreId GROUP BY g.GenreId ORDER BY TrackCount DESC",
+                    "user_input": "List genres by number of tracks"
+                },
+                {
+                    "query": "SELECT g.Name, SUM(t.Milliseconds)/3600000 as Hours FROM Genre g JOIN Track t ON g.GenreId = t.GenreId GROUP BY g.GenreId ORDER BY Hours DESC",
+                    "user_input": "Total hours of music by genre"
+                },
+                {
+                    "query": "SELECT g.Name, SUM(il.UnitPrice * il.Quantity) as Revenue FROM Genre g JOIN Track t ON g.GenreId = t.GenreId JOIN InvoiceLine il ON t.TrackId = il.TrackId GROUP BY g.GenreId ORDER BY Revenue DESC",
+                    "user_input": "Revenue generated by each genre"
+                }
+            ]
+        },
+        {
+            "separator": "table_6",
+            "name": "Invoice",
+            "schema": "CREATE TABLE Invoice (InvoiceId integer PRIMARY KEY, CustomerId integer, InvoiceDate timestamp without time zone, BillingAddress character varying(70), BillingCity character varying(40), BillingState character varying(40), BillingCountry character varying(40), BillingPostalCode character varying(10), Total numeric(10,2));",
+            "description": "This table stores invoices for customer purchases.",
+            "columns": [
+                {
+                    "name": "InvoiceId",
+                    "description": "unique identifier for invoices",
+                    "synonyms": [
+                        "invoice id"
+                    ]
+                },
+                {
+                    "name": "CustomerId",
+                    "description": "Id of the customer associated with the invoice",
+                    "synonyms": [
+                        "customer id"
+                    ]
+                },
+                {
+                    "name": "InvoiceDate",
+                    "description": "date when the invoice was generated",
+                    "synonyms": [
+                        "date",
+                        "purchase date"
+                    ]
+                },
+                {
+                    "name": "BillingAddress",
+                    "description": "billing address for the invoice",
+                    "synonyms": [
+                        "address"
+                    ]
+                },
+                {
+                    "name": "BillingCity",
+                    "description": "city for the billing address",
+                    "synonyms": [
+                        "city"
+                    ]
+                },
+                {
+                    "name": "BillingState",
+                    "description": "state or province for the billing address",
+                    "synonyms": [
+                        "state",
+                        "province"
+                    ]
+                },
+                {
+                    "name": "BillingCountry",
+                    "description": "country for the billing address",
+                    "synonyms": [
+                        "country"
+                    ]
+                },
+                {
+                    "name": "BillingPostalCode",
+                    "description": "postal code for the billing address",
+                    "synonyms": [
+                        "zip code",
+                        "postal code"
+                    ]
+                },
+                {
+                    "name": "Total",
+                    "description": "total amount of the invoice",
+                    "synonyms": [
+                        "amount",
+                        "price"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT SUM(Total) as TotalSales FROM Invoice WHERE strftime('%Y', InvoiceDate) = '2009'",
+                    "user_input": "Get total sales for 2009"
+                },
+                {
+                    "query": "SELECT strftime('%Y-%m', InvoiceDate) as Month, SUM(Total) as MonthlyRevenue FROM Invoice GROUP BY Month ORDER BY Month",
+                    "user_input": "Monthly revenue over time"
+                },
+                {
+                    "query": "SELECT BillingCountry, SUM(Total) as Revenue FROM Invoice GROUP BY BillingCountry ORDER BY Revenue DESC",
+                    "user_input": "Total revenue by country"
+                },
+                {
+                    "query": "SELECT i.InvoiceId, COUNT(il.InvoiceLineId) as Items FROM Invoice i JOIN InvoiceLine il ON i.InvoiceId = il.InvoiceId GROUP BY i.InvoiceId ORDER BY Items DESC LIMIT 1",
+                    "user_input": "Find the invoice with the most items"
+                }
+            ]
+        },
+        {
+            "separator": "table_7",
+            "name": "InvoiceLine",
+            "schema": "CREATE TABLE InvoiceLine (InvoiceLineId integer PRIMARY KEY, InvoiceId integer, TrackId integer, UnitPrice numeric(10,2), Quantity integer);",
+            "description": "This table stores individual line items for each invoice.",
+            "columns": [
+                {
+                    "name": "InvoiceLineId",
+                    "description": "unique identifier for invoice lines",
+                    "synonyms": [
+                        "line id"
+                    ]
+                },
+                {
+                    "name": "InvoiceId",
+                    "description": "Id of the invoice this line belongs to",
+                    "synonyms": [
+                        "invoice id"
+                    ]
+                },
+                {
+                    "name": "TrackId",
+                    "description": "Id of the track purchased",
+                    "synonyms": [
+                        "track id"
+                    ]
+                },
+                {
+                    "name": "UnitPrice",
+                    "description": "price of the track",
+                    "synonyms": [
+                        "price"
+                    ]
+                },
+                {
+                    "name": "Quantity",
+                    "description": "number of tracks purchased",
+                    "synonyms": [
+                        "amount",
+                        "count"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT t.Name, il.UnitPrice, il.Quantity FROM InvoiceLine il JOIN Track t ON il.TrackId = t.TrackId WHERE il.InvoiceId = 1",
+                    "user_input": "Get details of purchase in invoice 1"
+                },
+                {
+                    "query": "SELECT t.Name, SUM(il.Quantity) as TimesPurchased FROM InvoiceLine il JOIN Track t ON il.TrackId = t.TrackId GROUP BY il.TrackId ORDER BY TimesPurchased DESC LIMIT 10",
+                    "user_input": "Find the top 10 most purchased tracks"
+                },
+                {
+                    "query": "SELECT i.InvoiceId, SUM(il.UnitPrice * il.Quantity) as TotalAmount FROM InvoiceLine il JOIN Invoice i ON il.InvoiceId = i.InvoiceId GROUP BY i.InvoiceId HAVING TotalAmount > 10",
+                    "user_input": "Find invoices with total amount greater than $10"
+                },
+                {
+                    "query": "SELECT AVG(UnitPrice) as AveragePrice FROM InvoiceLine",
+                    "user_input": "Find the average unit price of purchased items"
+                }
+            ]
+        },
+        {
+            "separator": "table_8",
+            "name": "MediaType",
+            "schema": "CREATE TABLE MediaType (MediaTypeId integer PRIMARY KEY, Name character varying(120));",
+            "description": "This table stores types of media formats for music tracks.",
+            "columns": [
+                {
+                    "name": "MediaTypeId",
+                    "description": "unique identifier for media types",
+                    "synonyms": [
+                        "media type id"
+                    ]
+                },
+                {
+                    "name": "Name",
+                    "description": "name of the media type",
+                    "synonyms": [
+                        "media type name",
+                        "format"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT Name FROM MediaType",
+                    "user_input": "List all available media types"
+                },
+                {
+                    "query": "SELECT mt.Name, COUNT(t.TrackId) as TrackCount FROM MediaType mt JOIN Track t ON mt.MediaTypeId = t.MediaTypeId GROUP BY mt.MediaTypeId ORDER BY TrackCount DESC",
+                    "user_input": "Count tracks by media type"
+                },
+                {
+                    "query": "SELECT mt.Name, SUM(il.UnitPrice * il.Quantity) as Revenue FROM MediaType mt JOIN Track t ON mt.MediaTypeId = t.MediaTypeId JOIN InvoiceLine il ON t.TrackId = il.TrackId GROUP BY mt.MediaTypeId ORDER BY Revenue DESC",
+                    "user_input": "Revenue generated by each media type"
+                },
+                {
+                    "query": "SELECT mt.Name, SUM(t.Milliseconds)/3600000 as Hours FROM MediaType mt JOIN Track t ON mt.MediaTypeId = t.MediaTypeId GROUP BY mt.MediaTypeId ORDER BY Hours DESC",
+                    "user_input": "Total duration in hours of each media type"
+                }
+            ]
+        },
+        {
+            "separator": "table_9",
+            "name": "Playlist",
+            "schema": "CREATE TABLE Playlist (PlaylistId integer PRIMARY KEY, Name character varying(120));",
+            "description": "This table stores playlists of tracks.",
+            "columns": [
+                {
+                    "name": "PlaylistId",
+                    "description": "unique identifier for playlists",
+                    "synonyms": [
+                        "playlist id"
+                    ]
+                },
+                {
+                    "name": "Name",
+                    "description": "name of the playlist",
+                    "synonyms": [
+                        "playlist name"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT Name FROM Playlist WHERE PlaylistId = 1",
+                    "user_input": "Get the name of playlist with id 1"
+                },
+                {
+                    "query": "SELECT p.Name, COUNT(pt.TrackId) as TrackCount FROM Playlist p JOIN PlaylistTrack pt ON p.PlaylistId = pt.PlaylistId GROUP BY p.PlaylistId ORDER BY TrackCount DESC",
+                    "user_input": "List playlists by number of tracks"
+                },
+                {
+                    "query": "SELECT p.Name, SUM(t.Milliseconds)/60000 as Minutes FROM Playlist p JOIN PlaylistTrack pt ON p.PlaylistId = pt.PlaylistId JOIN Track t ON pt.TrackId = t.TrackId GROUP BY p.PlaylistId ORDER BY Minutes DESC",
+                    "user_input": "Find the total duration of each playlist in minutes"
+                },
+                {
+                    "query": "SELECT p.Name, COUNT(DISTINCT g.GenreId) as GenreCount FROM Playlist p JOIN PlaylistTrack pt ON p.PlaylistId = pt.PlaylistId JOIN Track t ON pt.TrackId = t.TrackId JOIN Genre g ON t.GenreId = g.GenreId GROUP BY p.PlaylistId ORDER BY GenreCount DESC",
+                    "user_input": "Count how many different genres are in each playlist"
+                }
+            ]
+        },
+        {
+            "separator": "table_10",
+            "name": "PlaylistTrack",
+            "schema": "CREATE TABLE PlaylistTrack (PlaylistId integer, TrackId integer, PRIMARY KEY (PlaylistId, TrackId));",
+            "description": "This junction table connects playlists and tracks in a many-to-many relationship.",
+            "columns": [
+                {
+                    "name": "PlaylistId",
+                    "description": "Id of the playlist",
+                    "synonyms": [
+                        "playlist id"
+                    ]
+                },
+                {
+                    "name": "TrackId",
+                    "description": "Id of the track in the playlist",
+                    "synonyms": [
+                        "track id"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT t.Name FROM Track t JOIN PlaylistTrack pt ON t.TrackId = pt.TrackId WHERE pt.PlaylistId = 1",
+                    "user_input": "Get all tracks in playlist 1"
+                },
+                {
+                    "query": "SELECT t.Name, COUNT(pt.PlaylistId) as PlaylistCount FROM Track t JOIN PlaylistTrack pt ON t.TrackId = pt.TrackId GROUP BY t.TrackId ORDER BY PlaylistCount DESC LIMIT 10",
+                    "user_input": "Find the top 10 tracks that appear in the most playlists"
+                },
+                {
+                    "query": "SELECT p1.Name, p2.Name, COUNT(pt1.TrackId) as CommonTracks FROM Playlist p1 JOIN PlaylistTrack pt1 ON p1.PlaylistId = pt1.PlaylistId JOIN PlaylistTrack pt2 ON pt1.TrackId = pt2.TrackId JOIN Playlist p2 ON pt2.PlaylistId = p2.PlaylistId WHERE p1.PlaylistId < p2.PlaylistId GROUP BY p1.PlaylistId, p2.PlaylistId ORDER BY CommonTracks DESC LIMIT 5",
+                    "user_input": "Find the top 5 playlist pairs that share the most tracks"
+                },
+                {
+                    "query": "SELECT g.Name, COUNT(pt.TrackId) as TracksInPlaylists FROM Genre g JOIN Track t ON g.GenreId = t.GenreId JOIN PlaylistTrack pt ON t.TrackId = pt.TrackId GROUP BY g.GenreId ORDER BY TracksInPlaylists DESC",
+                    "user_input": "Popularity of genres in playlists by track count"
+                }
+            ]
+        },
+        {
+            "separator": "table_11",
+            "name": "Track",
+            "schema": "CREATE TABLE Track (TrackId integer PRIMARY KEY, Name character varying(200), AlbumId integer, MediaTypeId integer, GenreId integer, Composer character varying(220), Milliseconds integer, Bytes integer, UnitPrice numeric(10,2));",
+            "description": "This table stores information about individual music tracks.",
+            "columns": [
+                {
+                    "name": "TrackId",
+                    "description": "unique identifier for tracks",
+                    "synonyms": [
+                        "track id"
+                    ]
+                },
+                {
+                    "name": "Name",
+                    "description": "name of the track",
+                    "synonyms": [
+                        "track name",
+                        "song title"
+                    ]
+                },
+                {
+                    "name": "AlbumId",
+                    "description": "Id of the album this track belongs to",
+                    "synonyms": [
+                        "album id"
+                    ]
+                },
+                {
+                    "name": "MediaTypeId",
+                    "description": "Id of the media type for this track",
+                    "synonyms": [
+                        "media type id",
+                        "format id"
+                    ]
+                },
+                {
+                    "name": "GenreId",
+                    "description": "Id of the genre this track belongs to",
+                    "synonyms": [
+                        "genre id"
+                    ]
+                },
+                {
+                    "name": "Composer",
+                    "description": "name of the composer of this track",
+                    "synonyms": [
+                        "writer",
+                        "songwriter"
+                    ]
+                },
+                {
+                    "name": "Milliseconds",
+                    "description": "length of the track in milliseconds",
+                    "synonyms": [
+                        "duration",
+                        "length"
+                    ]
+                },
+                {
+                    "name": "Bytes",
+                    "description": "size of the track file in bytes",
+                    "synonyms": [
+                        "file size"
+                    ]
+                },
+                {
+                    "name": "UnitPrice",
+                    "description": "price of the track",
+                    "synonyms": [
+                        "price",
+                        "cost"
+                    ]
+                }
+            ],
+            "sample_queries": [
+                {
+                    "query": "SELECT t.Name, a.Title FROM Track t JOIN Album a ON t.AlbumId = a.AlbumId WHERE t.Composer LIKE '%Johnson%'",
+                    "user_input": "Find all tracks composed by someone with Johnson in their name"
+                },
+                {
+                    "query": "SELECT Name, Milliseconds/60000 as Minutes FROM Track ORDER BY Milliseconds DESC LIMIT 10",
+                    "user_input": "Find the 10 longest tracks by duration"
+                },
+                {
+                    "query": "SELECT g.Name, AVG(t.UnitPrice) as AvgPrice FROM Track t JOIN Genre g ON t.GenreId = g.GenreId GROUP BY t.GenreId ORDER BY AvgPrice DESC",
+                    "user_input": "Average price of tracks by genre"
+                },
+                {
+                    "query": "SELECT Name, Composer, Milliseconds/60000 as Minutes FROM Track WHERE Composer IS NOT NULL AND Milliseconds = (SELECT MAX(Milliseconds) FROM Track)",
+                    "user_input": "Find the longest track and its composer"
+                }
+            ]
+        }
+    ]
+}

prompt.py

@@ -0,0 +1,14 @@
+from langchain_core.prompts import ChatPromptTemplate
+
+template = """Based on the table schema below, write a SQL query that would answer the user's question:
+{context}
+
+Question: {question}
+SQL Query:"""
+
+prompt = ChatPromptTemplate.from_messages(
+    [
+        ("system", "Given an input question, convert it to a SQL query. No pre-amble."),
+        ("human", template),
+    ]
+)

requirements.txt

@@ -0,0 +1,8 @@
+jq
+langchain 
+langchain-huggingface 
+sentence_transformers 
+langchain-google-genai 
+google-ai-generativelanguage==0.6.15
+langchain-community 
+langchain-chroma

utils.py

@@ -0,0 +1,22 @@
+import sqlite3
+import pandas as pd
+
+def execute_query_and_return_df(query):
+  """
+  Executes a SQL query and returns the results as a Pandas DataFrame.
+
+  Args:
+    query: The SQL query to execute.
+
+  Returns:
+    A Pandas DataFrame containing the query results.
+  """
+  conn = sqlite3.connect("/content/Chinook_Sqlite.sqlite")
+  try:
+    df = pd.read_sql_query(query, conn)
+    return df
+  except Exception as e:
+    print(f"Error executing query: {e}")
+    return None
+  finally:
+    conn.close()