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" #conn = pymssql.connect(server="127.0.0.1", port="33412", user="reader", password="passwd", database="db_name") # 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()