feat: Parent Docuement Retriever

app.py

@@ -11,24 +11,12 @@ from langchain_google_genai import (
     HarmBlockThreshold,
     HarmCategory,
 )
-
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.retrievers import ParentDocumentRetriever
+from langchain.storage import InMemoryStore
 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é."
+import tiktoken
 
 model = GoogleGenerativeAI(
     model=config.GOOGLE_CHAT_MODEL,
@@ -45,29 +33,45 @@ embedding = embeddings_model = GoogleGenerativeAIEmbeddings(
 
 vectordb = Chroma(persist_directory=config.STORAGE_PATH, embedding_function=embedding)
 
-retriever = SelfQueryRetriever.from_llm(
-    model,
-    vectordb,
-    document_content_description,
-    metadata_field_info,
-)
+## retriever
+
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, separators=["\n"])
+
+
+# The storage layer for the parent documents
+store = InMemoryStore()
+retriever = ParentDocumentRetriever(
+    vectorstore=vectordb,
+    docstore=store,
+    child_splitter=text_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)
+    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']
+                context += f"Article{i}:\n"+doc.page_content + f"\nSource: {source}" + separator
+        return context
 
     rag_chain = (
         {
-            "context": vectordb.as_retriever() | format_docs,
+            "context": retriever | format_docs,
             "question": RunnablePassthrough(),
         }
         | prompt
         | model
         | StrOutputParser()
     )
+    
 
     cl.user_session.set("rag_chain", rag_chain)
 

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

@@ -120,24 +120,7 @@ def process_docs(
         documents=splits,
         embedding=embeddings_model,
         persist_directory=persist_directory,
-    )
-
-    # 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",
-    )
-
-    print(f"Indexing stats: {index_result}")
+    ) 
 
     return doc_search