app.py

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()