Create app.py

app.py

@@ -0,0 +1,219 @@
+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.llms import AzureMLOnlineEndpoint
+import chromadb
+from langchain.chat_models.azureml_endpoint import ContentFormatterBase
+import json
+
+from langchain.chains import create_sql_query_chain
+
+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, 
+# this is often not enough and we need to specify it in the agent prompt using the suffix argument in the constructor.
+
+from langchain.agents import create_sql_agent, AgentType
+from langchain.agents.agent_toolkits import SQLDatabaseToolkit
+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"
+
+# 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 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)
+
+    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)
+
+custom_suffix = """
+Compose a query in the All_data table in the db database.
+Here is a description of each column:
+destn_area_name: The name of the destination area.
+destn_district_name: The name of the destination district.
+score: The score of the destination area.
+avg_days_todelr: The average number of days to deliver to the destination area.
+time_per: The time period of the data.
+orgn_area: The code of the origin area.
+orgn_dist: The code of the origin district.
+orgn_area_name: The name of the origin area.
+orgn_dist_name: The name of the origin district.
+destn_area: The code of the destination area.
+destn_dist: The code of the destination district.
+destn_area_name: The name of the destination area.
+destn_dist_name: The name of the destination district.
+prodt: The product type.
+rptg_start_date: The start date of the reporting period.
+rptg_end_date: The end date of the reporting period.
+mo: The month of the reporting period.
+pstl_qtr: The quarter of the Postal reporting period.
+pstl_yr: The year of the Postal reporting period.
+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
+                        )
+
+from langchain.prompts import PromptTemplate
+
+
+def build_sql_chain(llm, db):
+    
+    dialect = "Azure SQL"
+    table_info = "All_data"
+    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"}
+    fs = str(few_shots)
+                        
+
+    TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
+    Use the following format:
+
+    Question: "Question here"
+    SQLQuery: "SQL Query to run"
+    SQLResult: "Result of the SQLQuery"
+    Answer: "Final answer here"
+
+    Only use the following tables:
+
+    {table_info}.
+
+    Some examples of SQL queries that correspond to questions are:
+
+    \{"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"\}
+
+    Question: {input}"""
+
+    CUSTOM_PROMPT = PromptTemplate(
+        input_variables=["input", "table_info", "dialect"], template=TEMPLATE
+    )
+
+    # Set verbose=True to see the full prompt:
+    return create_sql_query_chain(llm=llm, db=db)
+
+sql_chain = build_sql_chain(llm, db)
+
+@cl.on_chat_start
+def main():
+    # Parse the command line arguments
+   # args = parse_arguments()
+
+
+    # activate/deactivate the streaming StdOut callback for LLMs
+    #callbacks = [StreamingStdOutCallbackHandler()]
+
+    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
+    m = message.content
+    res = sql_chain.invoke({"question": m})
+    # Call the chain asynchronously
+
+    print(res)
+    await cl.Message(content=res).send()