Spaces:
Configuration error
Configuration error
File size: 5,631 Bytes
05269f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | from src.internal.rag.retriever.retriever_types import (
DocumentType,
RetrievalResult,
)
from abc import ABC, abstractmethod
from typing import List
from langchain_core.documents import Document
from pathlib import Path
import logging
from langchain_community.document_loaders import (
PyMuPDFLoader,
Docx2txtLoader,
UnstructuredPowerPointLoader,
TextLoader
)
import asyncio
import hashlib
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class BaseDocumentLoader(ABC):
"""Abstract base class for document loaders"""
@abstractmethod
async def load_document(self, file_path: str) -> List[Document]:
"""Load document from file path"""
pass
@abstractmethod
def get_supported_extensions(self) -> List[str]:
"""Get supported file extensions"""
pass
class MultiFormatDocumentLoader(BaseDocumentLoader):
"""Document loader supporting multiple formats"""
def __init__(self):
self.loaders = {
DocumentType.PDF: self._load_pdf,
DocumentType.DOCX: self._load_docx,
DocumentType.PPT: self._load_ppt,
DocumentType.PPTX: self._load_pptx,
DocumentType.TXT: self._load_txt
}
async def load_document(self, file_path: str) -> List[Document]:
"""Load document based on file extension"""
try:
file_path = Path(file_path)
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
doc_type = self._get_document_type(file_path)
loader_func = self.loaders.get(doc_type)
if not loader_func:
raise ValueError(f"Unsupported file type: {doc_type}")
logger.info(f"Loading {doc_type} document: {file_path}")
documents = await loader_func(str(file_path))
for doc in documents:
doc.metadata.update({
"file_path": str(file_path),
"file_name": file_path.name,
"file_type": doc_type.value,
"file_size": file_path.stat().st_size,
"file_hash": self._calculate_file_hash(file_path)
})
return documents
except Exception as e:
logger.error(f"Error loading document {file_path}: {str(e)}")
raise
def get_supported_extensions(self) -> List[str]:
"""Get supported file extensions"""
return [".pdf", ".docx", ".ppt", ".pptx", ".txt"]
def _get_document_type(self, file_path: Path) -> DocumentType:
"""Determine document type from file extension"""
extension = file_path.suffix.lower()
mapping = {
".pdf": DocumentType.PDF,
".docx": DocumentType.DOCX,
".ppt": DocumentType.PPT,
".pptx": DocumentType.PPTX,
".txt": DocumentType.TXT
}
doc_type = mapping.get(extension)
if not doc_type:
raise ValueError(f"Unsupported file extension: {extension}")
return doc_type
def _calculate_file_hash(self, file_path: Path) -> str:
"""Calculate MD5 hash of file"""
hash_md5 = hashlib.md5()
with open(file_path, "rb") as f:
for chunk in iter(lambda: f.read(4096), b""):
hash_md5.update(chunk)
return hash_md5.hexdigest()
async def _load_pdf(self, file_path: str) -> List[Document]:
"""Load PDF document"""
try:
loader = PyMuPDFLoader(file_path)
documents = await asyncio.get_event_loop().run_in_executor(
None, loader.load
)
return documents
except Exception as e:
raise Exception(f"Error loading PDF: {str(e)}")
async def _load_docx(self, file_path: str) -> List[Document]:
"""Load DOCX document"""
try:
loader = Docx2txtLoader(file_path)
documents = await asyncio.get_event_loop().run_in_executor(
None, loader.load
)
return documents
except Exception as e:
raise Exception(f"Error loading DOCX: {str(e)}")
async def _load_ppt(self, file_path: str) -> List[Document]:
"""Load PPT document"""
try:
loader = UnstructuredPowerPointLoader(file_path)
documents = await asyncio.get_event_loop().run_in_executor(
None, loader.load
)
return documents
except Exception as e:
raise Exception(f"Error loading PPT: {str(e)}")
async def _load_pptx(self, file_path: str) -> List[Document]:
"""Load PPTX document"""
try:
loader = UnstructuredPowerPointLoader(file_path)
documents = await asyncio.get_event_loop().run_in_executor(
None, loader.load
)
return documents
except Exception as e:
raise Exception(f"Error loading PPTX: {str(e)}")
async def _load_txt(self, file_path: str) -> List[Document]:
"""Load TXT document"""
try:
loader = TextLoader(file_path)
documents = await asyncio.get_event_loop().run_in_executor(
None, loader.load
)
return documents
except Exception as e:
raise Exception(f"Error loading TXT: {str(e)}") |