HamzaHasan07/Retail_SQL_LLM

1from langchain_community.llms import GooglePalm2from langchain_google_genai import GoogleGenerativeAI3from langchain.utilities import SQLDatabase4from langchain_experimental.sql import SQLDatabaseChain5from langchain.prompts import SemanticSimilarityExampleSelector6from langchain.embeddings import HuggingFaceEmbeddings7from langchain.vectorstores import Chroma8from langchain.prompts import FewShotPromptTemplate9from langchain.chains.sql_database.prompt import PROMPT_SUFFIX, _mysql_prompt10from langchain.prompts.prompt import PromptTemplate11 12from few_shots import few_shots13 14import os15from dotenv import load_dotenv16load_dotenv()17 18 19def get_few_shot_db_chain():20    21    db_user = 'root';22    db_password = 'adminroot';23    db_host = 'localhost';24    db_name = 'Retail'25 26    db = SQLDatabase.from_uri(f"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}",27                          sample_rows_in_table_info=3)28    llm = GoogleGenerativeAI(model="models/text-bison-001", google_api_key=os.environ['GOOGLE_API_KEY'], temperature=0.2)29 30    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2');31 32    to_vectorize = [" ".join(values.values()) for values in few_shots]33    vector_db = Chroma.from_texts(to_vectorize, embedding=embeddings, metadatas=few_shots)34    example_selector = SemanticSimilarityExampleSelector(vectorstore = vector_db, k=2,);35    36    example_prompt = PromptTemplate(37    input_variables = ['Question', 'SQLQuery', 'SQLResult', 'Answer'],38    template = "\nQuestion: {Question}\nSQLQuery: {SQLQuery}\nSQLResult: {SQLResult}\nAnswer: {Answer}");39 40    few_shot_prompt = FewShotPromptTemplate(41    example_selector = example_selector,42    example_prompt = example_prompt,43    prefix = _mysql_prompt,44    suffix = PROMPT_SUFFIX,45    input_variables = ['input', 'table_info', 'top_k'])46 47    chain = SQLDatabaseChain.from_llm(llm, db, verbose=True, prompt=few_shot_prompt)48 49    return(chain)50 51if __name__ == "__main__":52    chain = get_few_shot_db_chain()53    print(chain.run("Identify the customer with the highest CLV"))54 55