Update main.py

main.py

@@ -5,24 +5,18 @@ import pandas as pd
 import numpy as np
 from typing import List
 from pathlib import Path
-from operator import itemgetter
+
 from langchain_openai import ChatOpenAI, OpenAI
 from langchain.schema.runnable.config import RunnableConfig
-#from langchain.schema import StrOutputParser
-from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
+from langchain.schema import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
 
 from langchain.agents import AgentExecutor
 from langchain.agents.agent_types import AgentType
 #from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent, create_csv_agent
-#from langchain_community.agent_toolkits import create_sql_agent
-from langchain.chains import create_sql_query_chain
+from langchain_community.agent_toolkits import create_sql_agent
 from langchain_community.utilities import SQLDatabase
 from sqlalchemy import create_engine
-from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
-
-
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough
 
 import chainlit as cl
 from chainlit.input_widget import TextInput, Select, Switch, Slider
@@ -80,28 +74,7 @@ def create_agent(filename: str):
         db = cl.user_session.get("db")
     # Create a SAL agent.
     #e.g agent_executor.invoke({"input": "Quel est le nombre de chargé d'affaires en agencement par entreprise?"})
-    #return create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=False)
-    execute_query = QuerySQLDataBaseTool(db=db)
-    write_query = create_sql_query_chain(llm, db)
-    
-    answer_prompt = PromptTemplate.from_template(
-        """Given the following user question, corresponding SQL query, and SQL result, answer the user question.
-
-    Question: {question}
-    SQL Query: {query}
-    SQL Result: {result}
-    Answer: """
-    )
-
-    chain = (
-        RunnablePassthrough.assign(query=write_query).assign(
-            result=itemgetter("query") | execute_query
-        )
-        | answer_prompt
-        | llm
-        | StrOutputParser()
-    )
-    cl.user_session.set("chain", chain)
+    return create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=False)
 
 def query_agent(agent, query):
     """
@@ -218,13 +191,12 @@ async def on_chat_start():
 async def on_message(message: cl.Message):
     await cl.Message(f"> SURVEYIA").send()
     agent = create_agent("./public/surveyia.csv")
-    chain_executor = cl.user_session.get("chain")
-    #cb = cl.AsyncLangchainCallbackHandler()
+    cb = cl.AsyncLangchainCallbackHandler()
     try:
         #res = await agent.acall("Réponds en langue française à la question suivante : " + message.content, callbacks=[cb])
-        res = await chain_executor.ainvoke({"question": "Réponds de la manière la plus complète et la plus intelligible, en langue française, à la question suivante : " + message.content + ". Réponds au format markdown ou au format tableau si le résultat nécessite l'affichage d'un tableau."})
+        res = await agent.ainvoke({"input": "Réponds de la manière la plus complète et la plus intelligible, en langue française, à la question suivante : " + message.content + ". Réponds au format markdown ou au format tableau si le résultat nécessite l'affichage d'un tableau."})
         #res = await agent.ainvoke("Réponds de la manière la plus complète et la plus intelligible, en langue française, à la question suivante : " + message.content + ". Réponds au format markdown ou au format tableau si le résultat nécessite l'affichage d'un tableau.")
-        await cl.Message(author="COPILOT",content=GoogleTranslator(source='auto', target='fr').translate(res)).send()
+        await cl.Message(author="COPILOT",content=GoogleTranslator(source='auto', target='fr').translate(res['output'])).send()
     except ValueError as e:
         res = str(e)
         resArray = res.split(":")