Update utils/DocsLoader.py

utils/DocsLoader.py

@@ -1,116 +1,28 @@
-# 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 = os.path.join("/tmp", "chunk_data")
-os.makedirs(CACHE_DIR, exist_ok=True)
 
 def load_and_chunk(url: str) -> list[Document]:
-    print(f"[Loader] URL: {url}")
+    print(url)
+    resp = requests.get(url)
+    if resp.status_code != 200:
+        raise HTTPException(400, "Could not download document")
 
-    # Try to get content length
-    try:
-        head = requests.head(url)
-        file_size = int(head.headers.get("Content-Length", 0))
-    except:
-        file_size = 0
-
-    USE_STREAM = file_size > 30 * 1024 * 1024  # >30MB
-
-    # Download content (streamed if large)
-    content = b""
-    hasher = hashlib.md5()
-
-    with requests.get(url, stream=USE_STREAM) as resp:
-        if resp.status_code != 200:
-            raise HTTPException(400, "Could not download document")
-
-        for chunk in resp.iter_content(chunk_size=8192):
-            hasher.update(chunk)
-            content += chunk
-
-    # Cache key based on content hash
-    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(content)
+            tmp.write(resp.content)
             tmp_path = tmp.name
         try:
             loader = PyMuPDFLoader(tmp_path)
@@ -118,9 +30,12 @@ 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(content)
+            tmp.write(resp.content)
             tmp_path = tmp.name
         try:
             loader = Docx2txtLoader(tmp_path)
@@ -129,19 +44,12 @@ def load_and_chunk(url: str) -> list[Document]:
             os.remove(tmp_path)
 
     elif "text/plain" in content_type or ".txt" in url_lower:
-        text = content.decode("utf-8", errors="ignore")
+        text = resp.content.decode("utf-8", errors="ignore")
         docs = [Document(page_content=text)]
 
     else:
         raise HTTPException(400, f"Unsupported document type: {content_type}")
 
-    # 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
+    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)