uspsoig-x/sqlchat

1            2from langchain.embeddings import HuggingFaceEmbeddings3#from langchain.vectorstores import FAISS4from langchain.schema import Document5 6#from langchain.vectorstores import Chroma 7from langchain.llms import AzureMLOnlineEndpoint8from langchain.chat_models.azureml_endpoint import ContentFormatterBase9import json10 11from langchain.chains import create_sql_query_chain12 13import chainlit as cl14 15 16from typing import Dict17 18 19# Now we can create the agent, adjusting the standard SQL Agent suffix to consider our use case. 20# Although the most straightforward way to handle this would be to include it just in the tool description, 21# this is often not enough and we need to specify it in the agent prompt using the suffix argument in the constructor.22 23from langchain.agents import create_sql_agent, AgentType24from langchain.agents.agent_toolkits import SQLDatabaseToolkit25from langchain.utilities import SQLDatabase26from langchain.chat_models import ChatOpenAI27 28import os29OPENAI_API_KEY = os.environ['OPENAI_API_KEY']30 31 32 33 34def create_agent():35   # conn_str = "mssql+pyodbc://" + SQL_USR_NM + ":" + PWD + "@" + SQL_HOST + "/" + SQL_TBL + "?driver=ODBC+Driver+18+for+SQL+Server"36 37    # Create the SQLDatabase object38    db = SQLDatabase.from_uri('sqlite:///spm.db')39    llm = ChatOpenAI(temperature=0.05, model="gpt-3.5-turbo-16k-0613")40    db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)41    return db_chain42 43    #toolkit = SQLDatabaseToolkit(db=db, llm=llm)44 45custom_suffix = """46Compose a query in the All_data table in the db database.47Here is a description of each column:48destn_area_name: The name of the destination area.49destn_district_name: The name of the destination district.50score: The score of the destination area.51avg_days_todelr: The average number of days to deliver to the destination area.52time_per: The time period of the data.53orgn_area: The code of the origin area.54orgn_dist: The code of the origin district.55orgn_area_name: The name of the origin area.56orgn_dist_name: The name of the origin district.57destn_area: The code of the destination area.58destn_dist: The code of the destination district.59destn_area_name: The name of the destination area.60destn_dist_name: The name of the destination district.61prodt: The product type.62rptg_start_date: The start date of the reporting period.63rptg_end_date: The end date of the reporting period.64mo: The month of the reporting period.65pstl_qtr: The quarter of the Postal reporting period.66pstl_yr: The year of the Postal reporting period.67score: The score of the destination area.68score_plus_1: The score of the destination area plus 1.69"""70 71#agent = create_sql_agent(llm=llm,72#                         toolkit=toolkit,73#                         verbose=False,74#                         agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,75#                         extra_tools=custom_tool_list,76#                         suffix=custom_suffix, 77#                         handle_parsing_errors=True78#                        )79 80from langchain.prompts import PromptTemplate81 82 83def build_sql_chain(llm, db):84    85    dialect = "Azure SQL"86    table_info = "All_data"87    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",88                        "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",89                        "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",90                        "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",91                        "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",92                        "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",93                        "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"}94    fs = str(few_shots)95                        96 97    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.98    Use the following format:99 100    Question: "Question here"101    SQLQuery: "SQL Query to run"102    SQLResult: "Result of the SQLQuery"103    Answer: "Final answer here"104 105    Only use the following tables:106 107    {table_info}.108 109    Some examples of SQL queries that correspond to questions are:110 111    \{"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",112     "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",113     "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",114     "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",115     "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",116     "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",117     "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"\}118 119    Question: {input}"""120 121    CUSTOM_PROMPT = PromptTemplate(122        input_variables=["input", "table_info", "dialect"], template=TEMPLATE123    )124 125    # Set verbose=True to see the full prompt:126    return create_sql_query_chain(llm=llm, db=db)127 128#from langchain.llms import OpenAI129from langchain_experimental.sql import SQLDatabaseChain130#sql_chain = build_sql_chain(llm, db)131 132 133@cl.on_chat_start134async def main():135    # Parse the command line arguments136   # args = parse_arguments()137    await cl.Message(content="Welcome to GeoData!").send()138    139    # activate/deactivate the streaming StdOut callback for LLMs140    #callbacks = [StreamingStdOutCallbackHandler()]141 142    #sql_chain = build_sql_chain(llm, db)143 144 145@cl.on_message146async def msg(message: str):147    # Retrieve the chain from the user session148    #sql_chain = cl.user_session.get("sql_chain")  # type: RetrievalQA149    agent = create_agent()150    m = message.content151    #res = sql_chain.invoke({"question": m})152    res = agent.run({"query": m})153    # Call the chain asynchronously154 155    print(res)156    await cl.Message(content=res).send()157