feat : Multiple documents & audit

app.py

@@ -29,6 +29,7 @@ def init_app():
         st.session_state["messages"] = []
         st.session_state["assistant"] = Rag()
         # st.session_state["data_dict"] = config['variables']
+        st.session_state["files"] = []
         st.session_state["prompt_system"] = config['prompt_system']
         st.session_state["chapters"] = config['chapters']
 

pages/documents.py

@@ -1,15 +1,19 @@
 import os
 import tempfile
 import streamlit as st
+from utils.document import audit_descriptif_pdf
 
 def read_and_save_file():
-    st.session_state["messages"] = []
-    st.session_state["user_input"] = ""
 
     for file in st.session_state["file_uploader"]:
         with tempfile.NamedTemporaryFile(delete=False) as tf:
             tf.write(file.getbuffer())
             file_path = tf.name
+            if not any(f["name"] == file.name for f in st.session_state["files"]):
+                st.session_state["files"].append({
+                    "name": file.name,
+                    "audit": audit_descriptif_pdf(file)["audit"]
+                })
 
         with st.session_state["ingestion_spinner"], st.spinner(f"Chargement {file.name}"):
             st.session_state["assistant"].ingest(file_path)
@@ -20,6 +24,9 @@ def read_and_save_file():
 def page():
     st.subheader("Charger vos documents")
 
+    if "files" not in st.session_state:
+        st.session_state["files"] = []
+
     # Custom CSS to hide default English labels
     # st.markdown(
     #     """
@@ -50,6 +57,32 @@ def page():
     )
 
 
+    for file in st.session_state["files"]:
+        st.markdown(f"#### {file['name']}")
+
+        audit = file["audit"]
+        st.markdown(
+            """
+            <table>
+                <tr><td>Nombre de pages</td><td>{}</td></tr>
+                <tr><td>Nombre d'images</td><td>{}</td></tr>
+                <tr><td>Nombre de liens</td><td>{}</td></tr>
+                <tr><td>Nombre de tableaux</td><td>{}</td></tr>
+                <tr><td>Nombre de tokens</td><td>{}</td></tr>
+                <tr><td>Nombre de mots</td><td>{}</td></tr>
+            </table>
+            """.format(
+                audit['number_of_pages'],
+                audit['number_of_images'],
+                audit['number_of_links'],
+                audit['number_of_tables'],
+                audit['number_of_tokens'],
+                audit['number_of_words']
+            ),
+            unsafe_allow_html=True
+        )
+
+
     st.session_state["ingestion_spinner"] = st.empty()
 
 page()

rag.py

@@ -10,7 +10,6 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.schema.runnable import RunnablePassthrough
 from langchain.prompts import PromptTemplate
 from langchain_community.vectorstores.utils import filter_complex_metadata
-from langchain_community.document_loaders.csv_loader import CSVLoader
 
 from util import getYamlConfig
 
@@ -24,6 +23,7 @@ class Rag:
     retriever = None
     chain = None
     readableModelName = ""
+    documents = []
 
     def __init__(self, vectore_store=None):
         
@@ -60,16 +60,18 @@ class Rag:
 
     def getDbFiles(self):
         return self.vector_store.getDocs()
-
+    
     def ingest(self, pdf_file_path: str):
         docs = PyPDFLoader(file_path=pdf_file_path).load()
-       
+
         chunks = self.text_splitter.split_documents(docs)
         chunks = filter_complex_metadata(chunks)
 
-        document_vector_store = FAISS.from_documents(chunks, self.embedding)
+        self.documents.extend(chunks)
+        self.document_vector_store = FAISS.from_documents(self.documents, self.embedding)
+
         
-        self.retriever = document_vector_store.as_retriever(
+        self.retriever = self.document_vector_store.as_retriever(
             search_type="similarity_score_threshold",
             search_kwargs={
                 "k": 3,
@@ -87,6 +89,8 @@ class Rag:
             documentContext = self.retriever.invoke(query)
 
 
+        print(documentContext)
+
         # Dictionnaire de base avec les variables principales
         chain_input = {
             "query": query,

requirements.txt

@@ -20,4 +20,5 @@ llamaapi
 pyyaml
 st_copy_to_clipboard
 faiss-gpu
-faiss-cpu
+faiss-cpu
+tiktoken

utils/document.py

@@ -0,0 +1,108 @@
+import pymupdf
+import tiktoken
+
+
+def count_tokens(input_string: str) -> int:
+    tokenizer = tiktoken.get_encoding("cl100k_base")
+    tokens = tokenizer.encode(input_string)
+    return len(tokens)
+
+
+def audit_descriptif_pdf(file,max_img_width=500) -> dict:
+    document = pymupdf.open(stream=file.read())
+
+    audit_dict_doc = {
+        "number_of_pages": len(document),
+        "number_of_images": 0,
+        "number_of_links": 0,
+        "number_of_tables": 0,
+        "number_of_tokens": 0,
+        "number_of_words": 0,
+        "key_words": []
+    }
+
+    doc_content = dict()
+
+    for page in document:
+
+        audit_dict_page = {}
+        page_content = {
+            "images": [],
+            "texte": "",
+            "liens": [],
+            "tableaux": []
+        }
+
+        #number of images
+        images = page.get_images()
+        number_images = len(images)
+        audit_dict_page["number_of_images"] = number_images
+        audit_dict_doc["number_of_images"] += number_images
+        
+        #get images
+        for _, img in enumerate(images):
+            xref = img[0]
+            base_image = document.extract_image(xref)
+
+            image_bytes = base_image["image"]
+            image_width = base_image["width"]
+            image_height = base_image["height"]
+            
+            # Adjust image size if it exceeds the maximum width
+            if image_width > max_img_width:
+                ratio = max_img_width / image_width
+                image_width = max_img_width
+                image_height = int(image_height * ratio)
+            
+            page_content["images"].append((image_bytes, image_width, image_height))
+
+
+        
+        #get links with uri
+        links = []
+        for link in page.get_links():
+            if link['kind'] == pymupdf.LINK_URI and 'uri' in link:
+                links.append({"uri": link["uri"], "page": page.number})
+        
+        page_content["liens"] = links
+
+        #number of links
+        number_links = len(links)
+        audit_dict_page["number_of_links"] = number_links
+        audit_dict_doc["number_of_links"] += number_links
+
+        #number of tables
+        tables = page.find_tables().tables
+        number_tables = len(tables)
+        for tab in tables:
+            page_content["tableaux"].append(tab.to_pandas())
+        audit_dict_page["number_of_tables"] = number_tables
+        audit_dict_doc["number_of_tables"] += number_tables
+
+        #number of tokens and words
+        text = page.get_text("text")
+        number_tokens = count_tokens(text)
+        number_words = len(text.split())
+
+        audit_dict_page["number_of_tokens"] = number_tokens
+        audit_dict_page["number_of_words"] = number_words
+
+        #get text
+        page_content["texte"] = text
+
+        audit_dict_doc["number_of_tokens"] += number_tokens
+        audit_dict_doc["number_of_words"] += number_words
+
+        audit_dict_doc[f"page_{page.number}"] = audit_dict_page
+
+        doc_content[f"page_{page.number}"] = page_content
+    
+
+    
+    #merge 2 dicts
+    global_audit = {
+        "audit": audit_dict_doc,
+        "content": doc_content
+    }
+
+    return global_audit