Update main.py

main.py

@@ -4,19 +4,14 @@ import bcrypt
 from typing import List
 from pathlib import Path
 from langchain_huggingface import HuggingFaceEmbeddings
-#from langchain_community.llms import HuggingFaceEndpoint
 from langchain_huggingface import HuggingFaceEndpoint
-#from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from langchain.prompts import ChatPromptTemplate
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.schema import StrOutputParser
-from langchain_community.document_loaders import (
-    PyMuPDFLoader,
-)
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import Chroma
 
-from langchain.indexes import SQLRecordManager, index
-from langchain.schema import Document
+from operator import itemgetter
+from pinecone import Pinecone
+
+from langchain.memory import ConversationBufferMemory
 from langchain.schema.runnable import Runnable, RunnablePassthrough, RunnableConfig
 from langchain.callbacks.base import BaseCallbackHandler
 
@@ -39,47 +34,6 @@ def auth_callback(username: str, password: str):
             identifier=ident + " : 🧑‍🎓 User Datapcc", metadata={"role": "user", "provider": "credentials"}
         )
 
-chunk_size = 1024
-chunk_overlap = 50
-
-embeddings_model = HuggingFaceEmbeddings()
-
-PDF_STORAGE_PATH = "./public/pdfs"
-
-
-def process_pdfs(pdf_storage_path: str):
-    pdf_directory = Path(pdf_storage_path)
-    docs = []  # type: List[Document]
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-
-    for pdf_path in pdf_directory.glob("*.pdf"):
-        loader = PyMuPDFLoader(str(pdf_path))
-        documents = loader.load()
-        docs += text_splitter.split_documents(documents)
-
-    doc_search = Chroma.from_documents(docs, embeddings_model)
-
-    namespace = "chromadb/my_documents"
-    record_manager = SQLRecordManager(
-        namespace, db_url="sqlite:///record_manager_cache.sql"
-    )
-    record_manager.create_schema()
-
-    index_result = index(
-        docs,
-        record_manager,
-        doc_search,
-        cleanup="incremental",
-        source_id_key="source",
-    )
-
-    print(f"Indexing stats: {index_result}")
-
-    return doc_search
-
-
-doc_search = process_pdfs(PDF_STORAGE_PATH)
-#model = ChatOpenAI(model_name="gpt-4", streaming=True)
 os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.environ['HUGGINGFACEHUB_API_TOKEN']
 repo_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
@@ -114,31 +68,38 @@ async def on_chat_start():
         await cl.Message(f"Vous pouvez requêter sur la thématique : {res.get('value')}").send()
         cl.user_session.set("selectRequest", res.get("value"))
         
-    template = """Answer the question based only on the following context:
-
+    memory = ConversationBufferMemory(return_messages=True)
+    template = """<s>[INST] 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. 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.
     {context}
-
-    Question: {question}
+    {question} [/INST] </s>
     """
-    prompt = ChatPromptTemplate.from_template(template)
-
-    def format_docs(docs):
-        return "\n\n".join([d.page_content for d in docs])
-
-    retriever = doc_search.as_retriever()
-
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                f"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 informations suivantes et du contexte suivant seulement et strictement.",
+            ),
+            MessagesPlaceholder(variable_name="history"),
+            ("human", "Contexte : {context}, réponds à la question suivante de la manière la plus pertinente, la plus exhaustive et la plus détaillée possible. {question}."),
+        ]
+    )
     runnable = (
-        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        RunnablePassthrough.assign(
+            history=RunnableLambda(memory.load_memory_variables) | itemgetter("history")
+        )
         | prompt
         | model
-        | StrOutputParser()
     )
 
+    cl.user_session.set("memory", memory)
     cl.user_session.set("runnable", runnable)
 
 
 @cl.on_message
 async def on_message(message: cl.Message):
+    runnable = cl.user_session.get("memory")
     runnable = cl.user_session.get("runnable")  # type: Runnable
     msg = cl.Message(content="")
 
@@ -159,21 +120,34 @@ async def on_message(message: cl.Message):
                 self.sources.add(source_page_pair)  # Add unique pairs to the set
 
         def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
-            cl.user_session.set("selectRequest","")
             if len(self.sources):
                 sources_text = "\n".join([f"{source}#page={page}" for source, page in self.sources])
                 self.msg.elements.append(
                     cl.Text(name="Sources", content=sources_text, display="inline")
                 )
 
+    os.environ['PINECONE_API_KEY'] = os.environ['PINECONE_API_KEY']
+    embeddings = HuggingFaceEmbeddings()
+    index_name = "all-venus"
+    pc = Pinecone(
+        api_key=os.environ['PINECONE_API_KEY']
+    )
+    index = pc.Index(index_name)
+    xq = embeddings.embed_query(message.content)
+    xc = index.query(vector=xq, filter={"categorie": {"$eq": "bibliographie-OPP-DGDIN"}},top_k=150, include_metadata=True)
+    context_p = ""
+    for result in xc['matches']:
+        context_p = context_p + result['metadata']['text']
+        
     async with cl.Step(type="run", name="QA Assistant"):
         async for chunk in runnable.astream(
-            cl.user_session.get("selectRequest"),
+            {"question": message.content, "context":context_p},
             config=RunnableConfig(callbacks=[
-                cl.LangchainCallbackHandler(),
-                PostMessageHandler(msg)
+                cl.AsyncLangchainCallbackHandler(stream_final_answer=True)
             ]),
         ):
             await msg.stream_token(chunk)
 
-    await msg.send()
+    await msg.send()
+    memory.chat_memory.add_user_message(message.content)
+    memory.chat_memory.add_ai_message(msg.content)