Update app.py

app.py

@@ -1,82 +1,105 @@
 import os
-import urllib.request
 import gradio as gr
-from llama_cpp import Llama
-from langchain.llms import llamacpp
-from huggingface_hub import login, hf_hub_download
-from dotenv import load_dotenv
+from dotenv import load_dotenv, find_dotenv
+from langchain.utilities.sql_database import SQLDatabase
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.tracers import ConsoleCallbackHandler
+from langchain_community.llms.llamacpp import LlamaCpp
+from huggingface_hub import login
+from langchain.globals import set_verbose
+set_verbose(True)
 
-MODEL_ID = "TheBloke/Llama-2-7b-Chat-GGUF" 
-MODEL_BASENAME = "llama-2-7b-chat.Q4_K_M.gguf"
-# MODEL_ID = "TheBloke/Wizard-Vicuna-7B-Uncensored-GGUF"
-# MODEL_BASENAME = "Wizard-Vicuna-7B-Uncensored.Q4_K_M.gguf"
-CONTEXT_WINDOW_SIZE = 8000
-MAX_NEW_TOKENS = 2000
-N_BATCH = 128
-# load_dotenv()
-os.getenv('hf_token')
-def load_quantized_model(model_id, model_basename):
-    try:
-        model_path = hf_hub_download(
-            repo_id=model_id,
-            filename=model_basename,
-            resume_download=True,
-            cache_dir="./models"
-        )
-        kwargs = {
-            'model_path': model_path,
-            'c_ctx': CONTEXT_WINDOW_SIZE,
-            'max_tokens': MAX_NEW_TOKENS,
-            'n_batch': N_BATCH
-        }
-        return llamacpp.LlamaCpp(**kwargs)
-    except TypeError:
-        return None
+# load_dotenv(find_dotenv(r"E:\AW\LLMs\.env"))
 
-def load_model(model_id, model_basename=None):
-    if ".gguf" in model_basename.lower():
-        llm = load_quantized_model(model_id, model_basename)
-        return llm
+def load_model(model_id):
+    if model_id == "gemini":
+        return ChatGoogleGenerativeAI(
+            model='gemini-pro',
+            google_api_key=login(os.environ["GOOGLE_API_KEY"]),
+            convert_system_message_to_human=True,
+            temperature=0.05,
+            verbose=True,
+        )
     else:
-        print("currently only .gguf models supported")   
-
+        print("only gemini supported aofn")
 
-
-def generate_text(prompt="Who is the CEO of Apple?"):
-    llm = load_model(MODEL_ID, MODEL_BASENAME)
-    output = llm(
-        prompt,
-        max_tokens=256,
-        temperature=0.1,
-        top_p=0.5,
-        echo=False,
-        stop=["#"],
+def chain(db, llm):
+    
+    def get_schema(_):
+        return db.get_table_info()
+    def run_query(query):
+        return db.run(query)
+    
+    template = """Based on the table schema below, write a MS SQL query that would answer the user's question:
+    {schema}
+    Question: {question}
+    Query:"""
+    
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "Given an input question, convert it to a MSSQL query. No pre-amble."),
+            ("human", template),
+        ]
     )
-    print(output)
-    return output
-    # output_text = output["choices"][0]["text"].strip()
-
-    # # Remove Prompt Echo from Generated Text
-    # cleaned_output_text = output_text.replace(prompt, "")
-    # return cleaned_output_text
-
-
-description = "Zephyr-beta"
+        
+    template = """Based on the table schema below, question, mssql query, and mssql response, write a natural language response:
+    {schema}
+        
+    Question: {question}
+    MS-SQL Query: {query}
+    MS-SQL Response: {response}"""
+        
+    prompt_response = ChatPromptTemplate.from_messages(
+        [
+            ("system", "Given an input question and MS-SQL response, convert it to a natural language answer. No pre-amble."),
+            ("human", template),
+        ]
+    )
+        
+    sql_response = (
+        RunnablePassthrough.assign(schema=get_schema)
+        | prompt
+        | llm.bind(stop=["\nSQLResult:"])
+        | StrOutputParser()
+    )
+    full_chain = (
+        RunnablePassthrough.assign(query=sql_response)
+        | RunnablePassthrough.assign(
+            schema=get_schema,
+            response=lambda x: db.run(x["query"]),
+        )
+        | prompt_response
+        | llm
+    )
+    
+    return full_chain
 
-examples = [
-    ["What is the capital of France?", "The capital of France is Paris."],
-    [
-        "Who wrote the novel 'Pride and Prejudice'?",
-        "The novel 'Pride and Prejudice' was written by Jane Austen.",
-    ],
-    ["What is the square root of 64?", "The square root of 64 is 8."],
-]
+def main():
+    gemini = load_model("gemini")
+    
+    path = r"OPPI_shift.db"       # \OPPI_down.db"
+    db1 = SQLDatabase.from_uri(f"sqlite:///{path}", include_tables=['ShiftDownTimeDetails'],sample_rows_in_table_info=0)
+    db2 = SQLDatabase.from_uri(f"sqlite:///{path}", include_tables=['ShiftProductionDetails'],sample_rows_in_table_info=0)
+    
+    down_chain = chain(db=db1, llm=gemini)
+    prod_chain = chain(db=db2, llm=gemini)
+    
+    def echo1(message, history):
+        ans = down_chain.invoke({"question":message}, config={"callbacks": [ConsoleCallbackHandler()]})
+        return str(ans)
+    
+    def echo2(message, history):
+        ans = prod_chain.invoke({"question":message}, config={"callbacks": [ConsoleCallbackHandler()]})
+        return str(ans)
+    
+    downtime = gr.ChatInterface(fn=echo1, title="SQL-Chatbot", description="Q/A on Downtime details table")
+    production = gr.ChatInterface(fn=echo2, title="SQL-Chatbot", description="Q/A on Production details table")
+    
+    demo = gr.TabbedInterface([downtime, production], ['ShiftDownTimeDetails', 'ShiftProductionDetails'])
+    demo.launch(debug=True)
 
-gradio_interface = gr.Interface(
-    fn=generate_text,
-    inputs="text",
-    outputs="text",
-    examples=examples,
-    title="Zephyr-B",
-)
-gradio_interface.launch(share=True)
+if __name__ == "__main__":
+    main()