Update main.py

main.py

@@ -20,6 +20,10 @@ from chainlit.input_widget import TextInput, Select, Switch, Slider
 from deep_translator import GoogleTranslator
 from IPython.display import display
 
+from literalai import LiteralClient
+literal_client = LiteralClient(api_key=os.getenv("LITERAL_API_KEY"))
+literal_client.instrument_openai()
+
 from surveycaa import surveyCaa
 
 @cl.password_auth_callback
@@ -39,6 +43,7 @@ def auth_callback(username: str, password: str):
             identifier=ident + " : 🧑‍🎓 User Datapcc", metadata={"role": "user", "provider": "credentials"}
         )
         
+@literal_client.step(type="tool")
 def create_agent(filename: str):
     """
     Create an agent that can access and use a large language model (LLM).
@@ -173,23 +178,24 @@ async def on_chat_start():
 
 @cl.on_message
 async def on_message(message: cl.Message):
-    await cl.Message(f"> SURVEYIA").send()
-    agent = create_agent("./public/surveyia.csv")
-    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("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:
-        res = str(e)
-        resArray = res.split(":")
-        ans = ''
-        if str(res).find('parsing') != -1:
-            for i in range(2,len(resArray)):
-                ans += resArray[i]
-            await cl.Message(author="COPILOT",content=ans.replace("`","")).send()
-        else:
-            await cl.Message(author="COPILOT",content="Reformulez votre requête, s'il vous plait 😃").send()
+    async with cl.Step(root=True, name="Réponse de l'Agent CSV", type="llm") as parent_step:
+        await cl.Message(f"> SURVEYIA").send()
+        agent = create_agent("./public/surveyia.csv")
+        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("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:
+            res = str(e)
+            resArray = res.split(":")
+            ans = ''
+            if str(res).find('parsing') != -1:
+                for i in range(2,len(resArray)):
+                    ans += resArray[i]
+                await cl.Message(author="COPILOT",content=ans.replace("`","")).send()
+            else:
+                await cl.Message(author="COPILOT",content="Reformulez votre requête, s'il vous plait 😃").send()
     # Query the agent.
     #response = query_agent(agent=agent, query=message.content)
     # Decode the response.