Update main.py

main.py

@@ -12,20 +12,16 @@ from operator import itemgetter
 from pinecone import Pinecone
 from langchain_pinecone import PineconeVectorStore
 from langchain_community.chat_message_histories import ChatMessageHistory
-#from langchain_google_community import GoogleSearchAPIWrapper
 from langchain.memory import ConversationBufferMemory
 from langchain.schema.runnable import Runnable, RunnablePassthrough, RunnableConfig, RunnableLambda
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import (
     StuffDocumentsChain, ConversationalRetrievalChain
 )
-#from langchain_core.tracers.context import tracing_v2_enabled
-#from langchain_core.tools import Tool
 
 import chainlit as cl
 from chainlit.input_widget import TextInput, Select, Switch, Slider
-#from chainlit.playground.config import add_llm_provider
-#from chainlit.playground.providers.langchain import LangchainGenericProvider
+from chainlit.types import ThreadDict
 
 from deep_translator import GoogleTranslator
 
@@ -51,25 +47,13 @@ def auth_callback(username: str, password: str):
             identifier=ident + " : 🧑‍🎓 User Datapcc", metadata={"role": "user", "provider": "credentials"}
         )
 
-@cl.step(type="llm")
+@cl.step(type="tool")
 async def LLModel():
     os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HUGGINGFACEHUB_API_TOKEN']
     repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
     llm = HuggingFaceEndpoint(
         repo_id=repo_id, max_new_tokens=5300, temperature=1.0, task="text2text-generation", streaming=True
     )
-    #add_llm_provider(
-    #    LangchainGenericProvider(
-            # It is important that the id of the provider matches the _llm_type
-    #        id=llm._llm_type,
-            # The name is not important. It will be displayed in the UI.
-    #        name="Mistral 8x7b Instruct",
-            # This should always be a Langchain llm instance (correctly configured)
-    #        llm=llm,
-            # If the LLM works with messages, set this to True
-    #        is_chat=True
-    #    )
-    #)
     return llm
     
 @cl.step(type="tool")
@@ -161,21 +145,24 @@ async def Search(input, categorie):
     results = [sources_text, verbatim_text, sources_offres]
     return results
     
+@cl.step(type="llm")
+async def setup_conversationalChain():
+    model = await LLModel()
+    retriever = await Retriever(res.get("name"))
+    qa = ConversationalRetrievalChain.from_llm(
+        model,
+        memory=memory,
+        chain_type="stuff",
+        return_source_documents=True,
+        verbose=False,
+        retriever=retriever
+    )
+    cl.user_session.set("runnable", qa)
+    cl.user_session.set("memory", memory)
+
 @cl.on_chat_start
 async def on_chat_start():
     await cl.Message(f"> REVIEWSTREAM").send()
-    #sources_videos = [
-    #    cl.Text(name="Videos", content="""<div style="display:inline-block;text-align:center;font-size:0.7rem;max-height:260px;width:33%;min-width:33%;max-width:33%;overflow:hidden"><a target="_blank" title="Alizé2 - Dimensionnement routier 1/2" href="https://clap.univ-eiffel.fr/permalink/v1261c4664a50dwedh6w/iframe/"><img src="https://clap.univ-eiffel.fr/thumb/v1261c4664a50dwedh6w/play/" width="100%" alt="Alizé2 - Dimensionnement routier 1/2"/><p>Alizé2 - Dimensionnement routier 1/2 Alizé2 - Dimensionnement routier 1/2</p></a></div>
-    #    <div style="display:inline-block;text-align:center;font-size:0.7rem;max-height:260px;width:33%;min-width:33%;max-width:33%;overflow:hidden"><a target="_blank" title="Alizé2 - Dimensionnement routier 1/2" href="https://clap.univ-eiffel.fr/permalink/v1261c4664a50dwedh6w/iframe/"><img src="https://clap.univ-eiffel.fr/thumb/v1261c4664a50dwedh6w/play/" width="100%" alt="Alizé2 - Dimensionnement routier 1/2"/><p>Alizé2 - Dimensionnement routier 1/2 Alizé2 - Dimensionnement routier 1/2</p></a></div>
-    #    <div style="display:inline-block;text-align:center;font-size:0.7rem;max-height:260px;width:33%;min-width:33%;max-width:33%;overflow:hidden"><a target="_blank" title="Alizé2 - Dimensionnement routier 1/2" href="https://clap.univ-eiffel.fr/permalink/v1261c4664a50dwedh6w/iframe/"><img src="https://clap.univ-eiffel.fr/thumb/v1261c4664a50dwedh6w/play/" width="100%" alt="Alizé2 - Dimensionnement routier 1/2"/><p>Alizé2 - Dimensionnement routier 1/2 Alizé2 - Dimensionnement routier 1/2</p></a></div>
-    #    <div style="display:inline-block;text-align:center;font-size:0.7rem;max-height:260px;width:33%;min-width:33%;max-width:33%;overflow:hidden"><a target="_blank" title="Alizé2 - Dimensionnement routier 1/2" href="https://clap.univ-eiffel.fr/permalink/v1261c4664a50dwedh6w/iframe/"><img src="https://clap.univ-eiffel.fr/thumb/v1261c4664a50dwedh6w/play/" width="100%" alt="Alizé2 - Dimensionnement routier 1/2"/><p>Alizé2 - Dimensionnement routier 1/2 Alizé2 - Dimensionnement routier 1/2</p></a></div>""",
-    #            display="inline")
-    #]
-    #await cl.Message(
-    #    content="Vidéos : ",
-    #    elements=sources_videos,
-    #).send()
-    
     res = await cl.AskActionMessage(
         content="<div style='width:100%;text-align:center'>Sélectionnez une source documentaire</div>",
         actions=[
@@ -217,20 +204,21 @@ async def on_chat_start():
     ########## Chain with streaming ##########
     message_history = ChatMessageHistory()
     memory = ConversationBufferMemory(memory_key="chat_history",output_key="answer",chat_memory=message_history,return_messages=True)
+    await setup_conversationalChain()
     
-    model = await LLModel()
-    retriever = await Retriever(res.get("name"))
-    qa = ConversationalRetrievalChain.from_llm(
-        model,
-        memory=memory,
-        chain_type="stuff",
-        return_source_documents=True,
-        verbose=False,
-        retriever=retriever
-    )
-    cl.user_session.set("runnable", qa)
+@cl.on_chat_resume
+async def on_chat_resume(thread: ThreadDict):
+    memory = ConversationBufferMemory(return_messages=True)
+    root_messages = [m for m in thread["steps"] if m["parentId"] == None]
+    for message in root_messages:
+        if message["type"] == "user_message":
+            memory.chat_memory.add_user_message(message["output"])
+        else:
+            memory.chat_memory.add_ai_message(message["output"])
+
     cl.user_session.set("memory", memory)
 
+    await setup_conversationalChain()
 
 @cl.on_message
 async def on_message(message: cl.Message):