Update utils/DocsLoader.py

utils/DocsLoader.py

@@ -1,28 +1,109 @@
+# import os
+# import tempfile
+# import requests
+
+# from fastapi import HTTPException
+# # from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
+# from langchain_community.document_loaders import PyMuPDFLoader, Docx2txtLoader
+# from langchain.text_splitter import RecursiveCharacterTextSplitter
+# from  langchain_text_splitters.sentence_transformers import SentenceTransformersTokenTextSplitter  # give better results but slow can use later for project
+# from langchain.schema import Document
+
+# MODEL_DIR = os.path.join("/tmp", "e5-large-v2")
+
+# 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)]
+
+#     else:
+#         raise HTTPException(400, f"Unsupported document type: {content_type}")
+
+#     # splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=120,separators=["\n\n", "\n", ".", " ", ""])
+#     splitter = SentenceTransformersTokenTextSplitter(model_name = MODEL_DIR,tokens_per_chunk=500, chunk_overlap=80)
+#     return splitter.split_documents(docs)
+
 import os
 import tempfile
 import requests
+import hashlib
+import pickle
 
 from fastapi import HTTPException
-# from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
 from langchain_community.document_loaders import PyMuPDFLoader, Docx2txtLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from  langchain_text_splitters.sentence_transformers import SentenceTransformersTokenTextSplitter  # give better results but slow can use later for project
 from langchain.schema import Document
+from langchain_text_splitters.sentence_transformers import SentenceTransformersTokenTextSplitter  # better but slower
 
 MODEL_DIR = os.path.join("/tmp", "e5-large-v2")
+CACHE_DIR = "/data/cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
 
 def load_and_chunk(url: str) -> list[Document]:
-    print(url)
-    resp = requests.get(url)
+    print(f"[Loader] Fetching: {url}")
+
+    # Get content length to decide whether to stream
+    head = requests.head(url)
+    file_size = int(head.headers.get("Content-Length", 0))
+    USE_STREAM = file_size > 30 * 1024 * 1024  # > 30MB
+
+    resp = requests.get(url, stream=USE_STREAM)
     if resp.status_code != 200:
         raise HTTPException(400, "Could not download document")
 
+    # Hash file content to create a unique cache key
+    hasher = hashlib.md5()
+    content = b""
+    for chunk in resp.iter_content(8192):
+        hasher.update(chunk)
+        content += chunk
+    content_hash = hasher.hexdigest()
+    cache_file = os.path.join(CACHE_DIR, f"{content_hash}.pkl")
+
+    # Return cached version if exists
+    if os.path.exists(cache_file):
+        print(f"✅ Loaded chunks from disk cache (key={content_hash})")
+        with open(cache_file, "rb") as f:
+            return pickle.load(f)
+
+    # Determine content type
     content_type = resp.headers.get("Content-Type", "").lower()
     url_lower = url.lower()
+    docs = []
 
     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.write(content)
             tmp_path = tmp.name
         try:
             loader = PyMuPDFLoader(tmp_path)
@@ -30,12 +111,9 @@ def load_and_chunk(url: str) -> list[Document]:
         finally:
             os.remove(tmp_path)
 
-    elif (
-        "application/vnd.openxmlformats-officedocument.wordprocessingml.document" in content_type
-        or ".docx" in url_lower
-    ):
+    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.write(content)
             tmp_path = tmp.name
         try:
             loader = Docx2txtLoader(tmp_path)
@@ -44,12 +122,20 @@ def load_and_chunk(url: str) -> list[Document]:
             os.remove(tmp_path)
 
     elif "text/plain" in content_type or ".txt" in url_lower:
-        text = resp.content.decode("utf-8", errors="ignore")
+        text = content.decode("utf-8", errors="ignore")
         docs = [Document(page_content=text)]
 
     else:
         raise HTTPException(400, f"Unsupported document type: {content_type}")
 
-    # splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=120,separators=["\n\n", "\n", ".", " ", ""])
-    splitter = SentenceTransformersTokenTextSplitter(model_name = MODEL_DIR,tokens_per_chunk=500, chunk_overlap=80)
-    return splitter.split_documents(docs)
+    # Split the docs
+    splitter = SentenceTransformersTokenTextSplitter(model_name=MODEL_DIR, tokens_per_chunk=500, chunk_overlap=80)
+    chunks = splitter.split_documents(docs)
+
+    # Save to cache
+    with open(cache_file, "wb") as f:
+        pickle.dump(chunks, f)
+    print(f"💾 Chunks cached to {cache_file}")
+
+    return chunks
+