|
import streamlit as st |
|
from langchain import HuggingFaceHub |
|
from langchain.prompts import ChatPromptTemplate |
|
from langchain_google_genai import ChatGoogleGenerativeAI, HarmBlockThreshold, HarmCategory |
|
from langchain.chains.question_answering import load_qa_chain |
|
from dotenv import load_dotenv |
|
import os |
|
import json |
|
import pandas as pd |
|
import numpy as np |
|
import sqlite3 |
|
|
|
|
|
load_dotenv() |
|
|
|
conn = sqlite3.connect('your_database.db') |
|
c = conn.cursor() |
|
|
|
|
|
def get_database_schema(): |
|
c.execute("SELECT name FROM sqlite_master WHERE type='table';") |
|
tables = c.fetchall() |
|
schema_dict = {} |
|
for table in tables: |
|
c.execute(f"PRAGMA table_info({table[0]})") |
|
schema = c.fetchall() |
|
schema_dict[table[0]] = schema |
|
return schema_dict |
|
|
|
|
|
def get_chain(): |
|
template = """ |
|
You are a SQLite query generator. |
|
return output executable query only or empyt text . |
|
output should be json with key "query" and tabel name with key "table_name" |
|
Based on user question, generate SQLite Query for given database Schema. |
|
Schema : \n{context}\n |
|
Question : {question} |
|
""" |
|
hf_model = HuggingFaceHub(repo_id="Mistralai/Mistralaichat", model_kwargs={"temperature": 0.4}) |
|
prompt=ChatPromptTemplate.from_template(template) |
|
chain=LLMChain(llm=hf_model.load_chat(),prompt=prompt) |
|
return chain |
|
|
|
def get_query(txt1): |
|
start_index = txt1.find('{') |
|
end_index = txt1.find('}') + 1 |
|
json_string = txt1[start_index:end_index] |
|
data = json.loads(json_string) |
|
|
|
return data |
|
|
|
|
|
def query_with_database(query, tablename=None): |
|
if tablename: |
|
c.execute(query) |
|
else: |
|
c.execute(query) |
|
query_output = c.fetchall() |
|
df = pd.DataFrame(query_output, columns=[desc[1] for desc in c.description]) |
|
st.table(df) |
|
|
|
|
|
def user_input(user_question): |
|
schema = get_database_schema() |
|
chain = get_chain() |
|
response = chain.run({"context": schema, "question": user_question}) |
|
generated_query = get_query(response) |
|
st.write("Generated Query : ", generated_query["query"]) |
|
query_with_database(generated_query["query"]) |
|
|
|
|
|
def main(): |
|
st.header("Chat With SQLite DATABASE - HUGGING FACE MISTRALAI/MISTRALAICHAT") |
|
with st.sidebar: |
|
st.write("Develop by : Rehan Shekh") |
|
st.write("Model : Hugging Face Mistralai/Mistralaichat") |
|
buttoclick=st.text_input(label="Ask from Database") |
|
st.write("Demo Database only have one table :- User_details") |
|
if buttoclick: |
|
|
|
user_input(str(buttoclick)) |
|
|
|
|
|
if __name__=="__main__": |
|
main() |