Update app.py

app.py

@@ -1,21 +1,10 @@
-import os
-os.environ['SENTENCE_TRANSFORMERS_HOME'] = './.cache'
-
-few_shots = {"What are the top 10 performing areas?": "SELECT TOP 10 destn_area_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_area_name ORDER BY AvgScore DESC",
-            "What are the worst 10 performing areas?": "SELECT TOP 10 destn_area_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_area_name ORDER BY AvgScore ASC",
-            "What districts have the highest volume of mail?": "SELECT TOP 10 destn_district_name, COUNT(*) AS Volume FROM All_data GROUP BY destn_district_name ORDER BY Volume DESC",
-            "What districts have the lowest volume of mail?": "SELECT TOP 10 destn_district_name, COUNT(*) AS Volume FROM All_data GROUP BY destn_district_name ORDER BY Volume ASC",
-            "What are the top 10 performing districts?": "SELECT TOP 10 destn_district_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_district_name ORDER BY AvgScore DESC",
-            "What are the worst 10 performing districts?": "SELECT TOP 10 destn_district_name, AVG(score) AS AvgScore FROM All_data GROUP BY destn_district_name ORDER BY AvgScore ASC",
-            "What districts gave the fastest delivery time?": "SELECT TOP 10 destn_district_name, AVG(avg_days_todelr) AS AvgDeliveryTime FROM All_data GROUP BY destn_district_name ORDER BY AvgDeliveryTime ASC"}
             
 from langchain.embeddings import HuggingFaceEmbeddings
 #from langchain.vectorstores import FAISS
 from langchain.schema import Document
 
-from langchain.vectorstores import Chroma 
+#from langchain.vectorstores import Chroma 
 from langchain.llms import AzureMLOnlineEndpoint
-import chromadb
 from langchain.chat_models.azureml_endpoint import ContentFormatterBase
 import json
 
@@ -26,26 +15,6 @@ import chainlit as cl
 
 from typing import Dict
 
-embeddings_model_name = 'sentence-transformers/msmarco-distilbert-base-tas-b'
-embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
-
-few_shot_docs = [Document(page_content=question, metadata={'sql_query': few_shots[question]}) for question in few_shots.keys()]
-vector_db = Chroma.from_documents(few_shot_docs, embeddings)
-retriever = vector_db.as_retriever()
-
-# Create custom tool and append it as a new tool in the create_sql_agent function:
-from langchain.agents.agent_toolkits import create_retriever_tool
-
-tool_description = """
-This tool will help answers questions about the USPS Service Performance Measurement (SPM) data.
-"""
-
-retriever_tool = create_retriever_tool(
-        retriever,
-        name='spm_chat',
-        description=tool_description
-    )
-custom_tool_list = [retriever_tool]
 
 # Now we can create the agent, adjusting the standard SQL Agent suffix to consider our use case. 
 # Although the most straightforward way to handle this would be to include it just in the tool description, 
@@ -57,60 +26,21 @@ from langchain.utilities import SQLDatabase
 from langchain.chat_models import ChatOpenAI
 
 import os
-
-PWD = os.environ['SQL_PWD']
-SQL_USR_NM = os.environ['SQL_USR_NM'] 
-SQL_HOST = os.environ['SQL_HOST']
-SQL_TBL = os.environ['SQL_TBL']
-
-conn_str = "mssql+pyodbc://" + SQL_USR_NM + ":" + PWD + "@" + SQL_HOST + "/" + SQL_TBL# + "#?driver=ODBC+Driver+18+for+SQL+Server"
-#conn = pymssql.connect(server="127.0.0.1", port="33412", user="reader", password="passwd", database="db_name")
+OPENAI_API_KEY = os.environ['OPENAI_API_KEY']
 
 
-# Create the SQLDatabase object
-db = SQLDatabase.from_uri(conn_str)
 
-model_name = os.environ['MODEL_NAME']
-endpoint_api_key = os.environ['ENDPOINT_API_KEY']
-endpoint_url = os.environ['ENDPOINT_URL']
 
-class CustomFormatter(ContentFormatterBase):
-    content_type = "application/json"
-    accepts = "application/json"
+def create_agent():
+   # conn_str = "mssql+pyodbc://" + SQL_USR_NM + ":" + PWD + "@" + SQL_HOST + "/" + SQL_TBL + "?driver=ODBC+Driver+18+for+SQL+Server"
 
-    def format_request_payload(self, prompt: str, model_kwargs: Dict) -> bytes:
-        print(model_kwargs)
-        input_str = json.dumps(
-            {
-            "input_data": {
-                "input_string": [
-                {
-                    "role": "user",
-                    "content": prompt
-                }
-                ],
-                "parameters": {
-                "temperature": 0.6,
-                "top_p": 0.9,
-                "max_new_tokens": 20000
-                }
-            }
-            }
-        )
-        return str.encode(input_str)
+    # Create the SQLDatabase object
+    db = SQLDatabase.from_uri('sqlite:///spm.db')
+    llm = ChatOpenAI(temperature=0.05, model="gpt-3.5-turbo-16k-0613")
+    db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)
+    return db_chain
 
-    def format_response_payload(self, output: bytes) -> str:
-        response_json = json.loads(output)
-        return response_json["output"]
-
-
-
-llm = AzureMLOnlineEndpoint(endpoint_name=model_name, 
-                endpoint_api_key=endpoint_api_key, 
-                endpoint_url=endpoint_url,
-                content_formatter = CustomFormatter())#,
-
-toolkit = SQLDatabaseToolkit(db=db, llm=llm)
+    #toolkit = SQLDatabaseToolkit(db=db, llm=llm)
 
 custom_suffix = """
 Compose a query in the All_data table in the db database.
@@ -138,14 +68,14 @@ score: The score of the destination area.
 score_plus_1: The score of the destination area plus 1.
 """
 
-agent = create_sql_agent(llm=llm,
-                         toolkit=toolkit,
-                         verbose=True,
-                       #  agent_type=AgentType.SELF_ASK_WITH_SEARCH,
-                         extra_tools=custom_tool_list,
-                         suffix=custom_suffix, 
-                         handle_parsing_errors=True
-                        )
+#agent = create_sql_agent(llm=llm,
+#                         toolkit=toolkit,
+#                         verbose=False,
+#                         agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+#                         extra_tools=custom_tool_list,
+#                         suffix=custom_suffix, 
+#                         handle_parsing_errors=True
+#                        )
 
 from langchain.prompts import PromptTemplate
 
@@ -195,27 +125,32 @@ def build_sql_chain(llm, db):
     # Set verbose=True to see the full prompt:
     return create_sql_query_chain(llm=llm, db=db)
 
-sql_chain = build_sql_chain(llm, db)
+#from langchain.llms import OpenAI
+from langchain_experimental.sql import SQLDatabaseChain
+#sql_chain = build_sql_chain(llm, db)
+
 
 @cl.on_chat_start
-def main():
+async def main():
     # Parse the command line arguments
    # args = parse_arguments()
-
-
+    await cl.Message(content="Welcome to GeoData!").send()
+    
     # activate/deactivate the streaming StdOut callback for LLMs
     #callbacks = [StreamingStdOutCallbackHandler()]
 
-    sql_chain = build_sql_chain(llm, db)
+    #sql_chain = build_sql_chain(llm, db)
 
 
 @cl.on_message
 async def msg(message: str):
     # Retrieve the chain from the user session
-   # sql_chain = cl.user_session.get("sql_chain")  # type: RetrievalQA
+    #sql_chain = cl.user_session.get("sql_chain")  # type: RetrievalQA
+    agent = create_agent()
     m = message.content
-    res = sql_chain.invoke({"question": m})
+    #res = sql_chain.invoke({"question": m})
+    res = agent.run({"query": m})
     # Call the chain asynchronously
 
     print(res)
-    await cl.Message(content=res).send()
+    await cl.Message(content=res).send()