Update main.py

main.py

@@ -19,7 +19,7 @@ from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import (
     StuffDocumentsChain, ConversationalRetrievalChain
 )
-from langchain_core.tracers.context import tracing_v2_enabled
+#from langchain_core.tracers.context import tracing_v2_enabled
 #from langchain_core.tools import Tool
 
 import chainlit as cl
@@ -31,8 +31,8 @@ from deep_translator import GoogleTranslator
 
 from datetime import timedelta
 
-from literalai import LiteralClient
-literal_client = LiteralClient(api_key=os.getenv("LITERAL_API_KEY"))
+from literalai import AsyncLiteralClient
+async_literal_client = AsyncLiteralClient(api_key=os.getenv("LITERAL_API_KEY"))
 
 @cl.password_auth_callback
 def auth_callback(username: str, password: str):
@@ -51,7 +51,8 @@ def auth_callback(username: str, password: str):
             identifier=ident + " : 🧑‍🎓 User Datapcc", metadata={"role": "user", "provider": "credentials"}
         )
 
-def LLModel():
+@cl.step(type="llm")
+async def LLModel():
     os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HUGGINGFACEHUB_API_TOKEN']
     repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
     llm = HuggingFaceEndpoint(
@@ -71,7 +72,8 @@ def LLModel():
     #)
     return llm
     
-def VectorDatabase(categorie):
+@cl.step(type="tool")
+async def VectorDatabase(categorie):
     if categorie == "bibliographie-OPP-DGDIN":
         index_name = "all-venus"
         embeddings = HuggingFaceEmbeddings()
@@ -93,8 +95,9 @@ def VectorDatabase(categorie):
     
     return vectorstore
 
-def Retriever(categorie):
-    vectorstore = VectorDatabase(categorie)
+@cl.step(type="retrieval")
+async def Retriever(categorie):
+    vectorstore = await VectorDatabase(categorie)
     if categorie == "bibliographie-OPP-DGDIN":
         retriever = vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": .7, "k": 150,"filter": {'categorie': {'$eq': categorie}}})
     elif categorie == "year":
@@ -105,8 +108,9 @@ def Retriever(categorie):
         retriever = vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={"score_threshold": .7, "k": 200,"filter": {"title": {"$eq": "videos-confinement-timeline"}}})
     return retriever
 
-def Search(input, categorie):
-    vectorstore = VectorDatabase(categorie)
+@cl.step(type="embedding")
+async def Search(input, categorie):
+    vectorstore = await VectorDatabase(categorie)
     results = []
     test = []
     sources_text = ""
@@ -214,13 +218,15 @@ async def on_chat_start():
     message_history = ChatMessageHistory()
     memory = ConversationBufferMemory(memory_key="chat_history",output_key="answer",chat_memory=message_history,return_messages=True)
     
+    model = await LLModel()
+    retriever = await Retriever(res.get("name"))
     qa = ConversationalRetrievalChain.from_llm(
-        LLModel(),
+        model,
         memory=memory,
         chain_type="stuff",
         return_source_documents=True,
         verbose=False,
-        retriever=Retriever(res.get("name"))
+        retriever=retriever
     )
     cl.user_session.set("runnable", qa)
     cl.user_session.set("memory", memory)
@@ -254,11 +260,10 @@ async def on_message(message: cl.Message):
                     cl.Text(name="Sources", content=sources_text, display="inline")
                 )
         
-    async with cl.Step(name="Réponse de Mistral 8x7b-Instruct", type="llm") as parent_step:
+    with literal_client.thread() as thread:
         cb = cl.AsyncLangchainCallbackHandler()
-        with tracing_v2_enabled(project_name="reviewStream"):
-            results = await runnable.acall("Contexte : Vous êtes un chercheur de l'enseignement supérieur et vous êtes doué pour faire des analyses d'articles de recherche sur les thématiques liées à la pédagogie, en fonction des critères définis ci-avant. En fonction des informations suivantes et du contexte suivant seulement et strictement, répondez en langue française strictement à la question ci-dessous à partir du contexte ci-dessous. En plus, tu créeras et tu afficheras 3 questions supplémentaires en relation avec le contexte initial, à chaque étape de la conversation. Tu écriras et tu afficheras les 3 questions supplémentaires en relation avec le contexte initial, avec un titrage de niveau 1 qui a pour titre \"Questions en relation avec le contexte : \". Lorsque cela est possible, cite les sources du contexte. Si vous ne pouvez pas répondre à la question sur la base des informations, dites que vous ne trouvez pas de réponse ou que vous ne parvenez pas à trouver de réponse. Essayez donc de comprendre en profondeur le contexte et répondez uniquement en vous basant sur les informations fournies. Ne générez pas de réponses non pertinentes. Question : " + message.content, callbacks=[cb])
-            answer = results["answer"]
+        results = await runnable.acall("Contexte : Vous êtes un chercheur de l'enseignement supérieur et vous êtes doué pour faire des analyses d'articles de recherche sur les thématiques liées à la pédagogie, en fonction des critères définis ci-avant. En fonction des informations suivantes et du contexte suivant seulement et strictement, répondez en langue française strictement à la question ci-dessous à partir du contexte ci-dessous. En plus, tu créeras et tu afficheras 3 questions supplémentaires en relation avec le contexte initial, à chaque étape de la conversation. Tu écriras et tu afficheras les 3 questions supplémentaires en relation avec le contexte initial, avec un titrage de niveau 1 qui a pour titre \"Questions en relation avec le contexte : \". Lorsque cela est possible, cite les sources du contexte. Si vous ne pouvez pas répondre à la question sur la base des informations, dites que vous ne trouvez pas de réponse ou que vous ne parvenez pas à trouver de réponse. Essayez donc de comprendre en profondeur le contexte et répondez uniquement en vous basant sur les informations fournies. Ne générez pas de réponses non pertinentes. Question : " + message.content, callbacks=[cb])
+        answer = results["answer"]
 
     await cl.Message(content=GoogleTranslator(source='auto', target='fr').translate(answer)).send()