Update app.py

app.py

@@ -3,13 +3,17 @@ import gradio as gr
 from dotenv import load_dotenv, find_dotenv
 from langchain.utilities.sql_database import SQLDatabase
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.chat_models.anthropic import ChatAnthropic
 from langchain_core.prompts import ChatPromptTemplate
+from langchain.agents import create_sql_agent, AgentType
 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 langchain.agents.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
 from huggingface_hub import login
 from langchain.globals import set_verbose
+from sqlalchemy import create_engine
+from prompts import agent_template, table_info
 set_verbose(True)
 
 # load_dotenv(find_dotenv(r".env"))
@@ -23,8 +27,16 @@ def load_model(model_id):
             temperature=0.05,
             verbose=True,
         )
+    elif model_name == "claude":
+        return ChatAnthropic(
+            model_name="claude-2",
+            anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"),
+            temperature=0.05,
+            streaming=True,
+            verbose=True,
+        )
     else:
-        print("only gemini supported aofn")
+        print("only gemini and claude supported aofn")
 
 def chain(db, llm):
     
@@ -79,10 +91,12 @@ def chain(db, llm):
 
 def main():
     gemini = load_model("gemini")
+    agent_llm = load_model("claude")
     
     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)
+    db3 = SQLDatabase.from_uri(f"sqlite:///{path}", include_tables=['ShiftDownTimeDetails','ShiftProductionDetails'],sample_rows_in_table_info=0)
     
     down_chain = chain(db=db1, llm=gemini)
     prod_chain = chain(db=db2, llm=gemini)
@@ -95,10 +109,29 @@ def main():
         ans = prod_chain.invoke({"question":message}, config={"callbacks": [ConsoleCallbackHandler()]})
         return str(ans)
     
+    prompt_agent = ChatPromptTemplate.from_messages(
+        [ 
+            ("system", "Given an input question, create a syntactically correct MS-SQL query to run, then look at the results of the query and return the answer in natural language. No Pre-amble."+agent_template),
+            ("human", "{question}"+table_info)
+        ]
+    )
+    sql_toolkit = SQLDatabaseToolkit(db=db3, llm=agent_llm)
+    agent = create_sql_agent(
+        toolkit=sql_toolkit,
+        llm=llm2,
+        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        verbose=True,
+        agent_executor_kwargs={"handle_parsing_errors":True}
+    )
+
+    def echo3(message, history):
+        answer = agent.invoke(prompt.format_prompt(question=message))
+        return answer['output']
+    
     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'])
+    agent = gr.ChatInterface(fn=echo3, title="SQL-Chatbot", description="General Chatbot with self-thinking capability, more robust to questions.")
+    demo = gr.TabbedInterface([agent, downtime, production], ['DB_bot-both tables','ShiftDownTimeDetails', 'ShiftProductionDetails'])
     demo.launch(debug=True)
 
 if __name__ == "__main__":