Update utils/DocsLoader.py

utils/DocsLoader.py

@@ -19,55 +19,66 @@ nltk.download("punkt_tab", download_dir=NLTK_PATH, quiet=True)
 
 MODEL_DIR = os.path.join("/tmp", "e5-large-v2")
 
+chunk_dict= {}
+    
+
 def load_and_chunk(url: str) -> list[Document]:
     print(url)
-    resp = requests.get(url)
-    if resp.status_code != 200:
-        raise HTTPException(400, "Could not download document")
-
-    content_type = resp.headers.get("Content-Type", "").lower()
-    url_lower = url.lower()
-
-    if "application/pdf" in content_type or ".pdf" in url_lower:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
-            tmp.write(resp.content)
-            tmp_path = tmp.name
-        try:
-            loader = PyMuPDFLoader(tmp_path)
-            docs = loader.load_and_split()
-        finally:
-            os.remove(tmp_path)
-
-    elif (
-        "application/vnd.openxmlformats-officedocument.wordprocessingml.document" in content_type
-        or ".docx" in url_lower
-    ):
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp:
-            tmp.write(resp.content)
-            tmp_path = tmp.name
-        try:
-            loader = Docx2txtLoader(tmp_path)
-            docs = loader.load_and_split()
-        finally:
-            os.remove(tmp_path)
-
-    elif "text/plain" in content_type or ".txt" in url_lower:
-        text = resp.content.decode("utf-8", errors="ignore")
-        docs = [Document(page_content=text)]
+    if(url not in chunk_dict):
+        print("processing new url")
+        resp = requests.get(url)
+        if resp.status_code != 200:
+            raise HTTPException(400, "Could not download document")
+    
+        content_type = resp.headers.get("Content-Type", "").lower()
+        url_lower = url.lower()
+    
+        if "application/pdf" in content_type or ".pdf" in url_lower:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+                tmp.write(resp.content)
+                tmp_path = tmp.name
+            try:
+                loader = PyMuPDFLoader(tmp_path)
+                docs = loader.load_and_split()
+            finally:
+                os.remove(tmp_path)
+    
+        elif (
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document" in content_type
+            or ".docx" in url_lower
+        ):
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp:
+                tmp.write(resp.content)
+                tmp_path = tmp.name
+            try:
+                loader = Docx2txtLoader(tmp_path)
+                docs = loader.load_and_split()
+            finally:
+                os.remove(tmp_path)
 
+        elif "text/plain" in content_type or ".txt" in url_lower:
+            text = resp.content.decode("utf-8", errors="ignore")
+            docs = [Document(page_content=text)]
+    
+        else:
+            raise HTTPException(400, f"Unsupported document type: {content_type}")
+        # # --- Step 1: Sentence split ---
+        # sentence_docs = []
+        # for doc in docs:
+        #     sentences = nltk.sent_tokenize(doc.page_content)
+        #     for sent in sentences:
+        #         if sent.strip():
+        #             sentence_docs.append(Document(page_content=sent, metadata=doc.metadata))
+        full_text = "\n".join([doc.page_content for doc in docs])  
+        
+        # splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
+        splitter = SentenceTransformersTokenTextSplitter(model_name = MODEL_DIR,tokens_per_chunk=512, chunk_overlap=90)
+        # return splitter.split_documents(docs)
+        # return splitter.split_documents([Document(page_content=full_text)])
+        chunk_dict[url] = splitter.create_documents([full_text])
+        # return splitter.create_documents([full_text])
+        return chunk_dict[url]
     else:
-        raise HTTPException(400, f"Unsupported document type: {content_type}")
-    # # --- Step 1: Sentence split ---
-    # sentence_docs = []
-    # for doc in docs:
-    #     sentences = nltk.sent_tokenize(doc.page_content)
-    #     for sent in sentences:
-    #         if sent.strip():
-    #             sentence_docs.append(Document(page_content=sent, metadata=doc.metadata))
-    full_text = "\n".join([doc.page_content for doc in docs])  
+        print("stored chunk")
+        return chunk_dict[url]
 
-    # splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
-    splitter = SentenceTransformersTokenTextSplitter(model_name = MODEL_DIR,tokens_per_chunk=512, chunk_overlap=80)
-    # return splitter.split_documents(docs)
-    # return splitter.split_documents([Document(page_content=full_text)])
-    return splitter.create_documents([full_text])