Update main.py

main.py

@@ -693,6 +693,7 @@ async def construction_FCS(romeListArray,settings):
 async def construction_NCS(romeListArray):
     context = await contexte(romeListArray)
     emploisST = cl.user_session.get("EmploiST")
+    memory = ConversationBufferMemory(return_messages=True)
     ### Mistral Completion ###
     client_llm = await IA()
     structure = str(modele('Note de composante sectorielle'))
@@ -710,11 +711,17 @@ async def construction_NCS(romeListArray):
     prompt = PromptTemplate(template=template, input_variables=["question","context"])
     #llm_chain = LLMChain(prompt=prompt, llm=client_llm)
     #completion_NCS = llm_chain.run({"question":question_p,"context":context_p}, callbacks=[StreamingStdOutCallbackHandler()])
-    chain = prompt | client_llm
+    chain = (
+        RunnablePassthrough.assign(
+            history=RunnableLambda(memory.load_memory_variables) | itemgetter("history")
+        )
+        | prompt | client_llm
+    )
     #completion_NCS = chain.invoke({"question":question_p,"context":context_p})
 
     msg = cl.Message(author="Datapcc : 🌐🌐🌐",content="")
-    async for chunk in chain.astream({"question":question_p,"context":context_p}):
+    async for chunk in chain.astream({"question":question_p,"context":context_p},
+                config=RunnableConfig(callbacks=[cl.AsyncLangchainCallbackHandler(stream_final_answer=True)]):
         await msg.stream_token(chunk)
     
     cl.user_session.set("NCS" + romeListArray[0], msg.content)
@@ -926,7 +933,7 @@ async def IA():
     repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
     llm = HuggingFaceEndpoint(
-        repo_id=repo_id, max_new_tokens=5000, temperature=0.7, streaming=True
+        repo_id=repo_id, max_new_tokens=5000, temperature=0.7, task="text2text-generation", streaming=True
     )
     return llm