Production-RAG-System/src/controllers/ProcessController.py at main · gamal1osama/Production-RAG-System · GitHub

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
from .BaseController import BaseController
from .ProjectController import ProjectController
from models import ProcessingEnums
from utils.PDFLoader import PDFLoader


from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os


class ProcessController(BaseController):

    def __init__(self, project_id: str):
        super().__init__()

        self.project_id = project_id
        self.project_path = ProjectController().get_project_path(project_id)

    def get_file_extension(self, file_id: str) -> str:
        return os.path.splitext(file_id)[-1]

    def get_file_loader(self, file_id: str):

        file_ext = self.get_file_extension(file_id)
        file_path = os.path.join(self.project_path, file_id)

        if not os.path.exists(file_path):
            print(f"Warning: File with id {file_id} not found in project {self.project_id}")
            return None

        if file_ext == ProcessingEnums.TXT.value:
            return TextLoader(file_path=file_path, encoding='utf-8') # we use utf-8 encoding because if the file contain arabic caracters

        elif file_ext == ProcessingEnums.PDF.value:
            return PDFLoader(file_path=file_path)

        return None

    def get_file_content(self, file_id: str) -> str:
        loader = self.get_file_loader(file_id)

        if loader is not None:
            return loader.load()

        return None

    def split_file_content(self, file_content: list,
                           chunk_size: int = 100, chunk_overlap: int = 20):

        text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size,
                                                       chunk_overlap=chunk_overlap,
                                                       separators=["\n"],
                                                       keep_separator=False)

        file_content_text = [content.page_content for content in file_content]
        file_content_metadata = [content.metadata for content in file_content]

        chunks = text_splitter.create_documents(file_content_text, metadatas=file_content_metadata)

        return chunks