Parthiban97/Chat_With_Multiple_SQL_Databases
1
1import streamlit as st2import os3import toml4from langchain_core.messages import AIMessage, HumanMessage5from langchain_core.prompts import ChatPromptTemplate6from langchain_core.runnables import RunnablePassthrough7from langchain_community.utilities import SQLDatabase8from langchain_core.output_parsers import StrOutputParser9from langchain_openai import ChatOpenAI10from langchain_groq import ChatGroq11#from dotenv import load_dotenv12 13#load_dotenv() 14 15# Function to update config.toml file16def update_secrets_file(user, data):17 secrets_file_path = ".streamlit/config.toml"18 secrets_data = {}19 20 # Load existing data from config.toml21 if os.path.exists(secrets_file_path):22 try:23 with open(secrets_file_path, "r") as file:24 secrets_data = toml.load(file)25 except toml.TomlDecodeError:26 secrets_data = {}27 28 # Update user-specific secrets data29 secrets_data[user] = data30 31 # Write updated data back to config.toml32 with open(secrets_file_path, "w") as file:33 toml.dump(secrets_data, file)34 35 36# Initialize database connections37def init_databases(user):38 secrets_file_path = ".streamlit/config.toml"39 secrets_data = {}40 if os.path.exists(secrets_file_path):41 with open(secrets_file_path, "r") as file:42 secrets_data = toml.load(file)43 44 user_data = secrets_data.get(user, {})45 46 db_connections = {}47 for database in user_data.get("Databases", "").split(','):48 database = database.strip()49 if database:50 db_uri = f"mysql+mysqlconnector://{user_data['User']}:{user_data['Password']}@{user_data['Host']}:{user_data['Port']}/{database}"51 db_connections[database] = SQLDatabase.from_uri(db_uri)52 return db_connections53 54 55# Function to get SQL chain56def get_sql_chain(dbs, llm):57 template = """58 You are a Senior and vastly experienced Data analyst at a company with around 20 years of experience. 59 You are interacting with a user who is asking you questions about the company's databases.60 Based on the table schemas below, write SQL queries that would answer the user's question. Take the conversation history into account.61 62 <SCHEMAS>{schemas}</SCHEMAS>63 64 Conversation History: {chat_history}65 66 Write the SQL queries for each relevant database, prefixed by the database name (e.g., DB1: SELECT * FROM ...; DB2: SELECT * FROM ...).67 Do not wrap the SQL queries in any other text, not even backticks.68 69 70 For example:71 72 Question: Which 3 artists have the most tracks?73 SQL Query: SELECT ArtistId, COUNT(*) as track_count FROM Track GROUP BY ArtistId ORDER BY track_count DESC LIMIT 3;74 75 Question: Name 10 artists76 SQL Query: SELECT Name FROM Artist LIMIT 10;77 78 Question:can you show SRK( a.k.a Shah Rukh Khan) movies where the imdb rating is greater than average rating of all movies and also show which movie had highest revenue in millions (INR)79 SQL Query: SELECT m.title, m.imdb_rating, f.revenue 80 FROM movies m 81 JOIN financials f ON m.movie_id = f.movie_id 82 JOIN movie_actor ma ON m.movie_id = ma.movie_id 83 JOIN actors a ON ma.actor_id = a.actor_id 84 WHERE a.name = 'Shah Rukh Khan' AND m.imdb_rating > (SELECT AVG(imdb_rating) FROM movies) 85 ORDER BY f.revenue DESC;86 87 Question: How many Van huesen black medium t shirts are available in stock?88 SQL Query: SELECT SUM(stock_quantity) FROM t_shirts WHERE brand = 'Van Huesen' AND color = 'Black' AND size = 'M';89 90 Question: How much is the price of the inventory for all small size t-shirts?91 SQL Query: SELECT SUM(price * stock_quantity) FROM t_shirts WHERE size = 'S';92 93 Question: If we have to sell all the Levi's T-shirts today with discounts applied, how much revenue our store will generate (post discounts)?94 SQL Query: SELECT SUM(a.total_amount * ((100 - COALESCE(discounts.pct_discount, 0)) / 100)) AS total_revenue FROM (SELECT SUM(price * stock_quantity) AS total_amount, t_shirt_id FROM t_shirts WHERE brand = 'Levi' GROUP BY t_shirt_id) a LEFT JOIN discounts ON a.t_shirt_id = discounts.t_shirt_id;95 96 Question: For each brand, find the total revenue generated from t-shirts with a discount applied, grouped by the discount percentage.97 SQL Query: SELECT brand, COALESCE(discounts.pct_discount, 0) AS discount_pct, SUM(t.price * t.stock_quantity * (1 - COALESCE(discounts.pct_discount, 0) / 100)) AS total_revenue FROM t_shirts t LEFT JOIN discounts ON t.t_shirt_id = discounts.t_shirt_id GROUP BY brand, COALESCE(discounts.pct_discount, 0);98 99 Question: Find the top 3 most popular colors for each brand, based on the total stock quantity.100 SQL Query: SELECT brand, color, SUM(stock_quantity) AS total_stock FROM t_shirts GROUP BY brand, color ORDER BY brand, total_stock DESC;101 102 Question: Calculate the average price per size for each brand, excluding sizes with less than 10 t-shirts in stock.103 SQL Query: SELECT brand, size, AVG(price) AS avg_price FROM t_shirts WHERE stock_quantity >= 10 GROUP BY brand, size HAVING COUNT(*) >= 10;104 105 Question: Find the brand and color combination with the highest total revenue, considering discounts.106 SQL Query: SELECT brand, color, SUM(t.price * t.stock_quantity * (1 - COALESCE(d.pct_discount, 0) / 100)) AS total_revenue FROM t_shirts t LEFT JOIN discounts d ON t.t_shirt_id = d.t_shirt_id GROUP BY brand, color ORDER BY total_revenue DESC LIMIT 1;107 108 Question: Create a view that shows the total stock quantity and revenue for each brand, size, and color combination.109 SQL Query: CREATE VIEW brand_size_color_stats AS SELECT brand, size, color, SUM(stock_quantity) AS total_stock, SUM(price * stock_quantity) AS total_revenue FROM t_shirts GROUP BY brand, size, color;110 111 Question: How much is the price of the inventory for all variants of t-shirts and group them by brands?112 SQL Query: SELECT brand, SUM(price * stock_quantity) FROM t_shirts GROUP BY brand;113 114 Question: List the total revenue of t-shirts of L size for all brands.115 SQL Query: SELECT brand, SUM(price * stock_quantity) AS total_revenue FROM t_shirts WHERE size = 'L' GROUP BY brand;116 117 Question: How many shirts are available in stock grouped by colours from each size and finally show me all brands?118 SQL Query: SELECT brand, color, size, SUM(stock_quantity) AS total_stock FROM t_shirts GROUP BY brand, color, size;119 120 Question: select all the movies with minimum and maximum release_year. Note that there can be more than one movies in min and max year hence output rows can be more than 2?121 SQL Query: select * from movies where release_year in (select min(release_year) from movies, select max(release_year) from movies);122 123 Question: Generate a yearly report for Croma India where there are two columns 1. Fiscal Year and 2. Total Gross Sales amount In that year from Croma124 SQL Query: select get_fiscal_year(date) as fiscal_year, sum(round(sold_quantity*g.gross_price,2)) as yearly_sales from fact_sales_monthly s join fact_gross_price g on g.fiscal_year=get_fiscal_year(s.date) and g.product_code=s.product_code where customer_code=90002002 group by get_fiscal_year(date) order by fiscal_year;125 126 Question: What is the total freight cost incurred by each customer in the month of May 2024?127 SQL Query: SELECT s.customer_name, SUM(f.freight_cost) AS total_freight_cost FROM gdb0041.sales_monthly s JOIN gdb056.freight_cost f ON s.customer_id = f.customer_id WHERE s.month = 'May 2024' GROUP BY s.customer_id;128 129 Question: Which market has the highest gross price sales in the last quarter?130 SQL Query: SELECT s.market, SUM(g.gross_price) AS total_gross_price FROM gdb041.sales_monthly s JOIN gdb056.gross_price g ON s.market_id = g.market_id WHERE s.quarter = 'Q2 2024' GROUP BY s.market_id ORDER BY total_gross_price DESC LIMIT 1;131 132 Question: What is the manufacturing cost of products sold in each region last year?133 SQL Query: SELECT s.region, SUM(m.manufacturing_cost) AS total_manufacturing_cost FROM gdb0041.sales_monthly s JOIN gdb056.manufacturing_cost m ON s.product_id = m.product_id WHERE s.year = 2023 GROUP BY s.region;134 135 Question: How many pre-invoice deductions were applied to each customer's sales in the last six months?136 SQL Query: SELECT s.customer_name, COUNT(p.pre_invoice_deduction_id) AS total_pre_invoice_deductions FROM gdb041.sales_monthly s JOIN gdb056.pre_invoice_deductions p ON s.sales_id = p.sales_id WHERE s.date BETWEEN DATE_SUB(NOW(), INTERVAL 6 MONTH) AND NOW() GROUP BY s.customer_id;137 138 Question: What are the post-invoice deductions for each product in the current year?139 SQL Query: SELECT f.product_name, SUM(p.amount) AS total_post_invoice_deductions FROM gdb0041.forecast_monthly f JOIN gdb056.post_invoice_deductions p ON f.product_id = p.product_id WHERE YEAR(f.date) = YEAR(NOW()) GROUP BY f.product_id;140 141 Question: How many movies did Sanjay Dutt act in, and which film gained the most revenue, ordered from highest to lowest revenue?142 SQL Query: SELECT m.title, SUM(b.revenue) AS total_revenue FROM movies m JOIN box_office b ON m.movie_id = b.movie_id JOIN movie_actor ma ON m.movie_id = ma.movie_id JOIN actors a ON ma.actor_id = a.actor_id WHERE a.name = 'Sanjay Dutt' GROUP BY m.title ORDER BY total_revenue DESC; 143 144 145 Your turn:146 147 Question: {question}148 SQL Queries:149 """150 151 prompt = ChatPromptTemplate.from_template(template)152 llm = llm153 154 def get_schema(_):155 schemas = {db_name: db.get_table_info() for db_name, db in dbs.items()}156 return schemas157 158 def parse_multi_line_queries(result):159 queries = {}160 lines = result.strip().split("\n")161 current_db = None162 current_query = []163 164 for line in lines:165 if ":" in line and not current_db: # Only split on colon for the database name166 current_db, query_start = line.split(":", 1)167 current_db = current_db.strip()168 current_query.append(query_start.strip())169 else:170 current_query.append(line.strip())171 172 if current_db:173 queries[current_db] = " ".join(current_query).strip()174 175 return queries176 177 return (178 RunnablePassthrough.assign(schemas=get_schema)179 | prompt180 | llm181 | StrOutputParser()182 | parse_multi_line_queries183 )184 185# Function to get response186def get_response(user_query, dbs, chat_history, llm):187 sql_chain = get_sql_chain(dbs, llm)188 189 template = """190 You are a Senior and vastly experienced Data analyst at a company with around 20 years of experience.191 You are interacting with a user who is asking you questions about the company's databases.192 Based on the table schemas below, question, sql queries, and sql responses, write an 193 accurate natural language response so that the end user can understand things194 and make sure do not include words like "Based on the SQL queries I ran". 195 Just provide only the answer with some text that the user expects.196 197 <SCHEMAS>{schemas}</SCHEMAS>198 Conversation History: {chat_history}199 SQL Queries: <SQL>{queries}</SQL>200 User question: {question}201 SQL Responses: {responses}202 """203 204 prompt = ChatPromptTemplate.from_template(template)205 llm = llm206 207 def run_queries(var):208 responses = {}209 for db_name, query in var["queries"].items():210 responses[db_name] = dbs[db_name].run(query)211 print(dbs.keys())212 print(var["queries"].keys())213 return responses214 215 chain = (216 RunnablePassthrough.assign(queries=sql_chain).assign(217 schemas=lambda _: {db_name: db.get_table_info() for db_name, db in dbs.items()},218 responses=run_queries,) # The comma at the end of the assign() method call is used to indicate that there may be more keyword arguments or method calls following it219 | prompt220 | llm221 | StrOutputParser()222 )223 224 return chain.invoke({225 "question": user_query,226 "chat_history": chat_history,227 })228 229# Streamlit app configuration230if "chat_history" not in st.session_state:231 st.session_state.chat_history = [232 AIMessage(content="Hello! I'm a SQL assistant. Ask me anything about your database."),233 ]234 235st.set_page_config(page_title="Chat with MySQL", page_icon="🛢️")236st.title("Chat with MySQL")237 238# Define model options239model_options_groq = [240 "llama-3.1-70b-versatile",241 "llama-3.1-8b-instant",242 "llama3-70b-8192",243 "llama3-8b-8192",244 "mixtral-8x7b-32768",245 "gemma2-9b-it"246]247 248model_options_openai = [249 "gpt-4o",250 "gpt-4o-mini",251 "gpt-4-turbo",252 "gpt-3.5-turbo-0125"253]254 255with st.sidebar:256 st.subheader("Settings")257 st.write("This is a simple chat application using MySQL. Connect to the database and start chatting.")258 259 if "db" not in st.session_state:260 st.session_state.user_id = st.text_input("User ID",placeholder="Enter any random numbers")261 st.session_state.Host = st.text_input("Host")262 st.session_state.Port = st.text_input("Port")263 st.session_state.User = st.text_input("User")264 st.session_state.Password = st.text_input("Password", type="password")265 st.session_state.Databases = st.text_input("Databases", placeholder="Enter DB's separated by (,)")266 st.session_state.openai_api_key = st.text_input("OpenAI API Key", type="password", help="Get your API key from [OpenAI Website](https://platform.openai.com/api-keys)")267 selected_model_openai = st.selectbox("Select any OpenAI Model", model_options_openai)268 st.session_state.groq_api_key = st.text_input("Groq API Key", type="password", help="Get your API key from [GROQ Console](https://console.groq.com/keys)")269 selected_model_groq = st.selectbox("Select any Groq Model", model_options_groq)270 st.info("Note: For interacting multiple databases and dealing with complex queries, GPT-4 Model is recommended for accurate results else proceed with Groq Model")271 272 273 os.environ["OPENAI_API_KEY"] = str(st.session_state.openai_api_key)274 275 276 if st.button("Connect"):277 with st.spinner("Connecting to databases..."):278 279 # Update config.toml with user-specific connection details280 update_secrets_file(st.session_state.user_id, {281 "Host": st.session_state.Host,282 "Port": st.session_state.Port,283 "User": st.session_state.User,284 "Password": st.session_state.Password,285 "Databases": st.session_state.Databases286 })287 288 dbs = init_databases(st.session_state.user_id)289 st.session_state.dbs = dbs290 291 if len(dbs) > 1:292 st.success(f"Connected to {len(dbs)} databases")293 else:294 st.success("Connected to database")295 296 297 298 if st.session_state.openai_api_key == "" and st.session_state.groq_api_key == "":299 st.error("Enter one API Key At least")300 elif st.session_state.openai_api_key: 301 st.session_state.llm = ChatOpenAI(model=selected_model_openai, api_key=st.session_state.openai_api_key)302 elif st.session_state.groq_api_key:303 st.session_state.llm = ChatGroq(model_name=selected_model_groq, temperature=0.5, api_key=st.session_state.groq_api_key)304 else:305 pass306 307# Display chat messages308for message in st.session_state.chat_history:309 if isinstance(message, AIMessage):310 with st.chat_message("AI"):311 st.markdown(message.content)312 elif isinstance(message, HumanMessage):313 with st.chat_message("Human"):314 st.markdown(message.content)315 316# Handle user input317user_query = st.chat_input("Type a message...")318if user_query is not None and user_query.strip() != "":319 st.session_state.chat_history.append(HumanMessage(content=user_query))320 321 with st.chat_message("Human"):322 st.markdown(user_query)323 324 with st.chat_message("AI"):325 response = get_response(user_query, st.session_state.dbs, st.session_state.chat_history, st.session_state.llm)326 st.markdown(response)327 328 st.session_state.chat_history.append(AIMessage(content=response))329