Update main.py

main.py

@@ -5,18 +5,24 @@ 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
+#from langchain.schema import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
 
 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_community.agent_toolkits import create_sql_agent
+from langchain.chains import create_sql_query_chain
 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
@@ -74,8 +80,28 @@ 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_pandas_dataframe_agent(llm, df, verbose=False, allow_dangerous_code=True, handle_parsing_errors=True, agent_type=AgentType.OPENAI_FUNCTIONS)
-    return create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=False)
+    #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)
 
 def query_agent(agent, query):
     """
@@ -192,10 +218,11 @@ async def on_chat_start():
 async def on_message(message: cl.Message):
     await cl.Message(f"> SURVEYIA").send()
     agent = create_agent("./public/surveyia.csv")
-    cb = cl.AsyncLangchainCallbackHandler()
+    chain_executor = cl.user_session.get("chain")
+    #cb = cl.AsyncLangchainCallbackHandler()
     try:
         #res = await agent.acall("Réponds en langue française à la question suivante : " + message.content, callbacks=[cb])
-        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 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("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['output'])).send()
     except ValueError as e: