Merge pull request #3 from data354/dev

app.py

@@ -1,9 +1,9 @@
 import chainlit as cl
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.chains.query_constructor.schema import AttributeInfo
-from langchain.retrievers.self_query.base import SelfQueryRetriever
-from langchain.schema import StrOutputParser
-from langchain.schema.runnable import Runnable, RunnableConfig, RunnablePassthrough
+from langchain.retrievers import ParentDocumentRetriever
+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 (
     GoogleGenerativeAI,
@@ -14,21 +14,7 @@ from langchain_google_genai import (
 
 import config
 from prompts import prompt
-
-metadata_field_info = [
-    AttributeInfo(
-        name="title",
-        description="Le titre de l'article",
-        type="string",
-    ),
-    AttributeInfo(
-        name="date",
-        description="Date de publication",
-        type="string",
-    ),
-    AttributeInfo(name="link", description="Source de l'article", type="string"),
-]
-document_content_description = "Articles sur l'actualité."
+from utils import PostMessageHandler, format_docs
 
 model = GoogleGenerativeAI(
     model=config.GOOGLE_CHAT_MODEL,
@@ -38,38 +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 = SelfQueryRetriever.from_llm(
-    model,
-    vectordb,
-    document_content_description,
-    metadata_field_info,
+## retriever
+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
+fs = LocalFileStore(config.STORAGE_PATH + "docstore")
+store = create_kv_docstore(fs)
+
+retriever = ParentDocumentRetriever(
+    vectorstore=vectorstore,
+    docstore=store,
+    child_splitter=child_splitter,
 )
 
 
 @cl.on_chat_start
 async def on_chat_start():
 
-    def format_docs(docs):
-        return "\n\n".join(doc.page_content for doc in docs)
-
-    rag_chain = (
-        {
-            "context": vectordb.as_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",
@@ -79,39 +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
-    msg = cl.Message(content="")
 
-    class PostMessageHandler(BaseCallbackHandler):
-        """
-        Callback handler for handling the retriever and LLM processes.
-        Used to post the sources of the retrieved documents as a Chainlit element.
-        """
-
-        def __init__(self, msg: cl.Message):
-            BaseCallbackHandler.__init__(self)
-            self.msg = msg
-            self.sources = []
-
-        def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
-            for d in documents:
-                source_doc = d.page_content + "\nSource: " + d.metadata["link"]
-                self.sources.append(source_doc)
-
-        def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
-            if len(self.sources):
-                # Display the reference docs with a Text widget
-                sources_element = [
-                    cl.Text(name=f"source_{idx+1}", content=content)
-                    for idx, content in enumerate(self.sources)
-                ]
-                source_names = [el.name for el in sources_element]
-                self.msg.elements += sources_element
-                self.msg.content += f"\nSources: {', '.join(source_names)}"
+    # 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> 
 
 """
 

requirements.txt

@@ -4,4 +4,5 @@ chainlit==1.0.500
 chromadb==0.4.24
 lark==1.1.9
 bs4==0.0.2
-selenium==4.19.0
+selenium==4.19.0
+tiktoken==0.1.1

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,45 +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"]
-    )
-    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)
+    # This text splitter is used to create the child documents
+    child_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n"])
 
-    doc_search = Chroma.from_documents(
-        documents=splits,
-        embedding=embeddings_model,
-        persist_directory=persist_directory,
+    # The vectorstore to use to index the child chunks
+    vectorstore = Chroma(
+        persist_directory=persist_directory + "vectorstore",
+        collection_name="full_documents",
+        embedding_function=embeddings_model,
     )
 
-    # Indexing data
-    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="link",
-    )
+    # The storage layer for the parent documents
+    fs = LocalFileStore(persist_directory + "docstore")
+    store = create_kv_docstore(fs)
 
-    print(f"Indexing stats: {index_result}")
+    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

@@ -0,0 +1,51 @@
+import chainlit as cl
+import tiktoken
+from langchain.callbacks.base import BaseCallbackHandler
+
+
+def format_docs(documents, max_context_size=100000, separator="\n\n"):
+    context = ""
+    encoder = tiktoken.get_encoding("cl100k_base")
+    i = 0
+    for doc in documents:
+        i += 1
+        if len(encoder.encode(context)) < max_context_size:
+            source = doc.metadata["link"]
+            title = doc.metadata["title"]
+            context += (
+                f"Article: {title}\n" + doc.page_content + f"\nSource: {source}" + separator
+            )
+    return context
+
+
+class PostMessageHandler(BaseCallbackHandler):
+    """
+    Callback handler for handling the retriever and LLM processes.
+    Used to post the sources of the retrieved documents as a Chainlit element.
+    """
+
+    def __init__(self, msg: cl.Message):
+        BaseCallbackHandler.__init__(self)
+        self.msg = msg
+        self.sources = []
+
+    def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
+        for d in documents:
+            source_doc = d.page_content + "\nSource: " + d.metadata["link"]
+            self.sources.append(source_doc)
+
+    def on_llm_end(self, response, *, run_id, parent_run_id, **kwargs):
+        if len(self.sources):
+            # Display the reference docs with a Text widget
+            sources_element = [
+                cl.Text(name=f"source_{idx+1}", content=content)
+                for idx, content in enumerate(self.sources)
+            ]
+            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