Perf(Parent Document Retriever): persist docs and vectorstore using LocalFileStore, update prompt, and refactor code

app.py

@@ -1,8 +1,8 @@
 import chainlit as cl
 from langchain.retrievers import ParentDocumentRetriever
-from langchain.schema import StrOutputParser
-from langchain.schema.runnable import Runnable, RunnableConfig, RunnablePassthrough
-from langchain.storage import InMemoryStore
+from langchain.schema.runnable import RunnableConfig
+from langchain.storage import LocalFileStore
+from langchain.storage._lc_store import create_kv_docstore
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores.chroma import Chroma
 from langchain_google_genai import (
@@ -24,39 +24,36 @@ model = GoogleGenerativeAI(
     },
 )  # type: ignore
 
-# Load vector database that was persisted earlier
-embedding = embeddings_model = GoogleGenerativeAIEmbeddings(
-    model="models/embedding-001", google_api_key=config.GOOGLE_API_KEY
+embeddings_model = GoogleGenerativeAIEmbeddings(
+    model=config.GOOGLE_EMBEDDING_MODEL
 )  # type: ignore
 
-vectordb = Chroma(persist_directory=config.STORAGE_PATH, embedding_function=embedding)
 
 ## retriever
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, separators=["\n"])
+child_splitter = RecursiveCharacterTextSplitter(chunk_size=500, separators=["\n"])
+
+# The vectorstore to use to index the child chunks
+vectorstore = Chroma(
+    persist_directory=config.STORAGE_PATH + "vectorstore",
+    collection_name="full_documents",
+    embedding_function=embeddings_model,
+)
 
 # The storage layer for the parent documents
-store = InMemoryStore()
+fs = LocalFileStore(config.STORAGE_PATH + "docstore")
+store = create_kv_docstore(fs)
+
 retriever = ParentDocumentRetriever(
-    vectorstore=vectordb,
+    vectorstore=vectorstore,
     docstore=store,
-    child_splitter=text_splitter,
+    child_splitter=child_splitter,
 )
 
 
 @cl.on_chat_start
 async def on_chat_start():
 
-    rag_chain = (
-        {
-            "context": retriever | format_docs,
-            "question": RunnablePassthrough(),
-        }
-        | prompt
-        | model
-        | StrOutputParser()
-    )
-
-    cl.user_session.set("rag_chain", rag_chain)
+    cl.user_session.set("retriever", retriever)
 
     msg = cl.Message(
         content=f"Vous pouvez poser vos questions sur les articles de SIKAFINANCE",
@@ -66,12 +63,20 @@ async def on_chat_start():
 
 @cl.on_message
 async def on_message(message: cl.Message):
-    runnable = cl.user_session.get("rag_chain")  # type: Runnable # type: ignore
+
+    # retriever = cl.user_session.get("retriever")
+
+    chain = prompt | model
+
     msg = cl.Message(content="")
 
     async with cl.Step(type="run", name="QA Assistant"):
-        async for chunk in runnable.astream(
-            message.content,
+
+        question = message.content
+        context = format_docs(retriever.get_relevant_documents(question))
+
+        async for chunk in chain.astream(
+            input={"context": context, "question": question},
             config=RunnableConfig(
                 callbacks=[cl.LangchainCallbackHandler(), PostMessageHandler(msg)]
             ),

config.py

@@ -3,7 +3,7 @@ import os
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 GOOGLE_CHAT_MODEL = "gemini-pro"
 GOOGLE_EMBEDDING_MODEL = "models/embedding-001"
-STORAGE_PATH = "data/chroma/"
+STORAGE_PATH = "./data/"
 HIISTORY_FILE = "./data/qa_history.txt"
 
 NUM_DAYS_PAST = 30

prompts.py

@@ -1,11 +1,27 @@
 from langchain.prompts import ChatPromptTemplate
 
 template = """
-Répondez à la question en vous basant uniquement sur le contexte suivant:
+Vous êtes un assistant de recherche économique et financière, spécialement conçu pour répondre aux questions liées à l'économie et à la finance et pour aider à l'informations et la prise de décisions financières. Votre rôle consiste à analyser les articles et rapports d'actualité économique et financière qui vous sera fournis dans le contexte et à répondre de manière adequate aux questions spécifiques des utilisateurs. Lorsque vous répondez aux questions :
+- Pour des questions d'ordre générales (ex: "Quelle est l'actualité du jour?") : Lisez attentivement tous les articles et résumez les points\évènements clés en mentionnant les dates de publications.
+- Pour des questions spécifiques (ex: "Quelle est la tendance du marché boursier aujourd'hui?") : Recherchez les informations spécifiques à la question dans les articles.
 
-{context}
+-N'hésitez pas à utiliser vos connaissances et votre bon sens pour répondre aux questions.
 
-Question : {question}
+- Basez vos réponses sur les articles d'actualité fournis. Citez directement les parties pertinentes de ces documents pour étayer vos réponses.
+- Citez clairement les références, y compris les titres des articles, les dates de publication et tout autre détail pertinent, afin de vous assurer que les informations peuvent être facilement vérifiées et retracées jusqu'aux sources originales.
+
+- Si la question sort du cadre des documents fournis ou si vous ne trouvez pas d'informations pertinentes, indiquez poliment que la réponse ne peut être déterminée sur la base des sources disponibles. Suggérez de consulter d'autres articles d'actualité financière ou des bases de données pour obtenir une réponse complète, le cas échéant.
+- Insistez sur l'exactitude et la fiabilité de vos réponses, en comprenant la nature critique de votre aide dans les processus de prise de décision financière.
+- Répondez aux utilisateurs dans la langue de leur question. Si la question est en français, votre réponse doit être en français. Si la question est en anglais, votre réponse doit être en anglais.
+- Pour des question en relative à la date veuillez considerer qu'aujourd'hui est le Jeudi 11/04/2024. Par exemple pour repondre à une question sur l'actualité du jour, vous devez effectuer une comparaison entre les date de publications des articles et celle d'aujourdui pour filtrer sur les articles puis retourner les informations pertinantes.
+
+<contexte> 
+``{context}`` 
+</contexte> 
+
+<question> 
+{question} 
+</question> 
 
 """
 

scrape_data.py

@@ -2,7 +2,9 @@ import os
 from datetime import date, timedelta
 
 import bs4
-from langchain.indexes import SQLRecordManager, index
+from langchain.retrievers import ParentDocumentRetriever
+from langchain.storage import LocalFileStore
+from langchain.storage._lc_store import create_kv_docstore
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores.chroma import Chroma
 from langchain_community.document_loaders import WebBaseLoader
@@ -81,7 +83,7 @@ def set_metadata(documents, metadatas):
 
 
 def process_docs(
-    articles, persist_directory, embeddings_model, chunk_size=1000, chunk_overlap=100
+    articles, persist_directory, embeddings_model, chunk_size=500, chunk_overlap=0
 ):
     """
     #Scrap all articles urls content and save on a vector DB
@@ -105,28 +107,33 @@ def process_docs(
     # Update metadata: add title,
     set_metadata(documents=docs, metadatas=articles)
 
-    print("Successfully loaded to document")
+    # print("Successfully loaded to document")
 
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n"]
+    # This text splitter is used to create the child documents
+    child_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n"])
+
+    # The vectorstore to use to index the child chunks
+    vectorstore = Chroma(
+        persist_directory=persist_directory + "vectorstore",
+        collection_name="full_documents",
+        embedding_function=embeddings_model,
     )
-    splits = text_splitter.split_documents(docs)
 
-    # Create the storage path if it doesn't exist
-    if not os.path.exists(persist_directory):
-        os.makedirs(persist_directory)
+    # The storage layer for the parent documents
+    fs = LocalFileStore(persist_directory + "docstore")
+    store = create_kv_docstore(fs)
 
-    doc_search = Chroma.from_documents(
-        documents=splits,
-        embedding=embeddings_model,
-        persist_directory=persist_directory,
+    retriever = ParentDocumentRetriever(
+        vectorstore=vectorstore,
+        docstore=store,
+        child_splitter=child_splitter,
     )
 
-    return doc_search
+    retriever.add_documents(docs, ids=None)
+    print(len(docs), " documents added")
 
 
 if __name__ == "__main__":
 
     data = scrap_articles(DATA_URL, num_days_past=config.NUM_DAYS_PAST)
-    vectordb = process_docs(data, config.STORAGE_PATH, embeddings_model)
-    ret = vectordb.as_retriever()
+    process_docs(data, config.STORAGE_PATH, embeddings_model)

utils.py

@@ -11,8 +11,9 @@ def format_docs(documents, max_context_size=100000, separator="\n\n"):
         i += 1
         if len(encoder.encode(context)) < max_context_size:
             source = doc.metadata["link"]
+            title = doc.metadata["title"]
             context += (
-                f"Article{i}:\n" + doc.page_content + f"\nSource: {source}" + separator
+                f"Article: {title}\n" + doc.page_content + f"\nSource: {source}" + separator
             )
     return context
 
@@ -43,3 +44,8 @@ class PostMessageHandler(BaseCallbackHandler):
             source_names = [el.name for el in sources_element]
             self.msg.elements += sources_element
             self.msg.content += f"\nSources: {', '.join(source_names)}"
+
+    def clean_text(text):
+        tx = text.replace("Tweet","")
+        tx = tx.replace("\n\n\n\n\n\n\n\n\n","")
+        return tx