import langchain
from dotenv import load_dotenv
from langchain.agents import initialize_agent, AgentType
from langchain.chat_models import ChatOpenAI
from datetime import timedelta, datetime
import chainlit as cl
from utils.custom_tools import CustomTrinoListTable, CustomTrinoTableSchema, CustomTrinoSqlQuery, CustomTrinoSqlCheck, CustomTrinoTableJoin

# 加载.env文件中的环境变量
load_dotenv()
langchain.debug = True
today = (datetime.now()).strftime("%Y%m%d")
yeaterday = (datetime.now() - timedelta(days=1)).strftime("%Y%m%d")
today_d = (datetime.now()).strftime("%Y-%m-%d")
yeaterday_d = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d")
custom_prefix = f""" You are an agent designed to interact with a Trino SQL database.
Please do not answer others questions,use Chinese to answer questions,think it step by step. 
NOTE: 
data rule:
date string format YYYYMMDD. for example, today is {today}, then yeaterday is {yeaterday}, and so on.
date string format YYYY-MM-DD. for example, today is {today_d}, then yeaterday is {yeaterday_d}, and so on.
column dt is table data version string, not a business date.
sql rule:
sql generation not need to end with ; only sql itself.
get answer from single sql after plan logics step by step. 
You have access to the following tools:"""


@cl.on_chat_start
async def start():
    # Initialize model
    model = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, verbose=False, streaming=True)

    custom_tool_list = [CustomTrinoListTable(), CustomTrinoTableSchema(), CustomTrinoSqlQuery(), CustomTrinoSqlCheck(), CustomTrinoTableJoin()]

    agent_executor = initialize_agent(
        custom_tool_list,
        llm=model,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        max_iterations=6,
        agent_kwargs={"prefix": custom_prefix},
        handle_parsing_errors="Check your output and make sure it conforms"
    )
    cl.user_session.set("agent", agent_executor)
    # Send the initial message
    elements = [
        cl.Text(name="提问：", content="计算订单明细表的dt为昨天,销售日期为本月的总销售数量", display="inline"),
        cl.Text(name="我能生成SQL脚本：", content=f"SELECT SUM(num) AS total_sales_quantity FROM gjdw.dw_sale_tr_goods_dt WHERE dt = '{today}' AND dates >= '2023-10-01' AND dates <= '2023-10-31'", display="inline", language="SQL"),
        cl.Text(name="最终结果：", content="订单明细表的dt为昨天，销售日期为本月的总销售数量是0。", display="inline"),
    ]
    content = "Hi,我是 Trino SQL Agent ,我能帮助你查询trino数据库。您可以向我提问，例如："
    await cl.Message(content=content, elements=elements).send()


@cl.on_message
async def main(message: cl.Message):
    agent = cl.user_session.get("agent")  # type: #AgentExecutor
    cb = cl.LangchainCallbackHandler(stream_final_answer=True)
    print(message)
    await cl.make_async(agent.run)(message, callbacks=[cb])

# def ask(input: str) -> str:
#     print("-- Serving request for input: %s" % input)
#     try:
#         response = agent_executor.run(input)
#     except Exception as e:
#         response = str(e)
#         if response.startswith("Could not parse LLM output: `"):
#             response = response.removeprefix("Could not parse LLM output: `").removesuffix("`")
#     return response


# agent_executor.run(" table gjdw.dw_sale_tr_goods_dt has column named bill_code ?")
# ask("计算订单明细表的dt为昨天,销售日期为本月初的总销售数量")

langchain.agents.initialize_agent langchain.chat_models.ChatOpenAI
[((371, 384), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (382, 384), False, 'from dotenv import load_dotenv\n'), ((1351, 1439), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)', 'verbose': '(False)', 'streaming': '(True)'}), "(model_name='gpt-3.5-turbo', temperature=0, verbose=False,\n    streaming=True)\n", (1361, 1439), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1603, 1848), 'langchain.agents.initialize_agent', 'initialize_agent', (['custom_tool_list'], {'llm': 'model', 'agent': 'AgentType.ZERO_SHOT_REACT_DESCRIPTION', 'verbose': '(True)', 'max_iterations': '(6)', 'agent_kwargs': "{'prefix': custom_prefix}", 'handle_parsing_errors': '"""Check your output and make sure it conforms"""'}), "(custom_tool_list, llm=model, agent=AgentType.\n    ZERO_SHOT_REACT_DESCRIPTION, verbose=True, max_iterations=6,\n    agent_kwargs={'prefix': custom_prefix}, handle_parsing_errors=\n    'Check your output and make sure it conforms')\n", (1619, 1848), False, 'from langchain.agents import initialize_agent, AgentType\n'), ((1901, 1945), 'chainlit.user_session.set', 'cl.user_session.set', (['"""agent"""', 'agent_executor'], {}), "('agent', agent_executor)\n", (1920, 1945), True, 'import chainlit as cl\n'), ((2597, 2625), 'chainlit.user_session.get', 'cl.user_session.get', (['"""agent"""'], {}), "('agent')\n", (2616, 2625), True, 'import chainlit as cl\n'), ((2659, 2712), 'chainlit.LangchainCallbackHandler', 'cl.LangchainCallbackHandler', ([], {'stream_final_answer': '(True)'}), '(stream_final_answer=True)\n', (2686, 2712), True, 'import chainlit as cl\n'), ((417, 431), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (429, 431), False, 'from datetime import timedelta, datetime\n'), ((531, 545), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (543, 545), False, 'from datetime import timedelta, datetime\n'), ((1461, 1483), 'utils.custom_tools.CustomTrinoListTable', 'CustomTrinoListTable', ([], {}), '()\n', (1481, 1483), False, 'from utils.custom_tools import CustomTrinoListTable, CustomTrinoTableSchema, CustomTrinoSqlQuery, CustomTrinoSqlCheck, CustomTrinoTableJoin\n'), ((1485, 1509), 'utils.custom_tools.CustomTrinoTableSchema', 'CustomTrinoTableSchema', ([], {}), '()\n', (1507, 1509), False, 'from utils.custom_tools import CustomTrinoListTable, CustomTrinoTableSchema, CustomTrinoSqlQuery, CustomTrinoSqlCheck, CustomTrinoTableJoin\n'), ((1511, 1532), 'utils.custom_tools.CustomTrinoSqlQuery', 'CustomTrinoSqlQuery', ([], {}), '()\n', (1530, 1532), False, 'from utils.custom_tools import CustomTrinoListTable, CustomTrinoTableSchema, CustomTrinoSqlQuery, CustomTrinoSqlCheck, CustomTrinoTableJoin\n'), ((1534, 1555), 'utils.custom_tools.CustomTrinoSqlCheck', 'CustomTrinoSqlCheck', ([], {}), '()\n', (1553, 1555), False, 'from utils.custom_tools import CustomTrinoListTable, CustomTrinoTableSchema, CustomTrinoSqlQuery, CustomTrinoSqlCheck, CustomTrinoTableJoin\n'), ((1557, 1579), 'utils.custom_tools.CustomTrinoTableJoin', 'CustomTrinoTableJoin', ([], {}), '()\n', (1577, 1579), False, 'from utils.custom_tools import CustomTrinoListTable, CustomTrinoTableSchema, CustomTrinoSqlQuery, CustomTrinoSqlCheck, CustomTrinoTableJoin\n'), ((2002, 2078), 'chainlit.Text', 'cl.Text', ([], {'name': '"""提问："""', 'content': '"""计算订单明细表的dt为昨天,销售日期为本月的总销售数量"""', 'display': '"""inline"""'}), "(name='提问：', content='计算订单明细表的dt为昨天,销售日期为本月的总销售数量', display='inline')\n", (2009, 2078), True, 'import chainlit as cl\n'), ((2088, 2313), 'chainlit.Text', 'cl.Text', ([], {'name': '"""我能生成SQL脚本："""', 'content': 'f"""SELECT SUM(num) AS total_sales_quantity FROM gjdw.dw_sale_tr_goods_dt WHERE dt = \'{today}\' AND dates >= \'2023-10-01\' AND dates <= \'2023-10-31\'"""', 'display': '"""inline"""', 'language': '"""SQL"""'}), '(name=\'我能生成SQL脚本：\', content=\n    f"SELECT SUM(num) AS total_sales_quantity FROM gjdw.dw_sale_tr_goods_dt WHERE dt = \'{today}\' AND dates >= \'2023-10-01\' AND dates <= \'2023-10-31\'"\n    , display=\'inline\', language=\'SQL\')\n', (2095, 2313), True, 'import chainlit as cl\n'), ((2313, 2392), 'chainlit.Text', 'cl.Text', ([], {'name': '"""最终结果："""', 'content': '"""订单明细表的dt为昨天，销售日期为本月的总销售数量是0。"""', 'display': '"""inline"""'}), "(name='最终结果：', content='订单明细表的dt为昨天，销售日期为本月的总销售数量是0。', display='inline')\n", (2320, 2392), True, 'import chainlit as cl\n'), ((465, 479), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (477, 479), False, 'from datetime import timedelta, datetime\n'), ((482, 499), 'datetime.timedelta', 'timedelta', ([], {'days': '(1)'}), '(days=1)\n', (491, 499), False, 'from datetime import timedelta, datetime\n'), ((583, 597), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (595, 597), False, 'from datetime import timedelta, datetime\n'), ((600, 617), 'datetime.timedelta', 'timedelta', ([], {'days': '(1)'}), '(days=1)\n', (609, 617), False, 'from datetime import timedelta, datetime\n'), ((2742, 2766), 'chainlit.make_async', 'cl.make_async', (['agent.run'], {}), '(agent.run)\n', (2755, 2766), True, 'import chainlit as cl\n'), ((2477, 2523), 'chainlit.Message', 'cl.Message', ([], {'content': 'content', 'elements': 'elements'}), '(content=content, elements=elements)\n', (2487, 2523), True, 'import chainlit as cl\n')]