_utils/splitters/Splitter_class.py

from _utils.bubble_integrations.obter_arquivo import get_pdf_from_bubble
from _utils.handle_files import return_document_list_with_llama_parser
from _utils.splitters.splitter_util import combine_documents_without_losing_pagination
from setup.easy_imports import PyPDFLoader, RecursiveCharacterTextSplitter, Document
from typing import Any, List, Dict, Tuple, Optional, cast
from _utils.models.gerar_relatorio import (
    DocumentChunk,
)
import uuid


class Splitter:
    def __init__(
        self,
        chunk_size,
        chunk_overlap,
    ):
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size, chunk_overlap=chunk_overlap
        )
        self.chunk_metadata = {}  # Store chunk metadata for tracing

    async def load_and_split_document(
        self, pdf_path: str, should_use_llama_parse: bool, isBubble: bool
    ):
        """Load PDF and split into chunks with metadata"""
        # loader = PyPDFLoader(pdf_path)
        # if not pages:
        #     pages = get_pdf_from_bubble(
        #         pdf_path
        #     )  # Gera uma lista de objetos Document, sendo cada item da lista referente a UMA PÁGINA inteira do PDF.

        initial_chunks: List[str] = []

        if isBubble:
            print("\nPEGANDO PDF DO BUBBLE")
            pages = await get_pdf_from_bubble(pdf_path, should_use_llama_parse)
            page_boundaries, combined_text = (
                combine_documents_without_losing_pagination(pages)
            )
            initial_chunks = initial_chunks + self.text_splitter.split_text(
                combined_text
            )
        else:
            if should_use_llama_parse:
                print("\nENVIANDO PDFS PARA LLAMA PARSE")
                pages = await return_document_list_with_llama_parser(pdf_path)
                page_boundaries, combined_text = (
                    combine_documents_without_losing_pagination(pages)
                )
                initial_chunks = initial_chunks + self.text_splitter.split_text(
                    combined_text
                )
            else:
                print("\nCOMEÇANDO LEITURA DO PDF")
                pages = PyPDFLoader(pdf_path).load()
                print("TERMINOU LEITURA DO PDF")
                page_boundaries, combined_text = (
                    combine_documents_without_losing_pagination(pages)
                )

                initial_chunks = initial_chunks + self.text_splitter.split_text(
                    combined_text
                )

        chunks: List[DocumentChunk] = []
        char_count = 0

        # for page in pages:
        #     text = page.page_content
        #     page_chunks = self.text_splitter.split_text(
        #         text
        #     )  # Quebra o item que é um Document de UMA PÁGINA inteira em um lista onde cada item é referente a um chunk, que são pedaços menores do que uma página.
        text_char = 0
        print("\nQUEBRANDO PDF EM CHUNKS ORGANIZADOS")
        for chunk in initial_chunks:
            chunk_id = str(uuid.uuid4())
            start_char = text_char + 1
            end_char = start_char + len(chunk)
            text_char = end_char

            if should_use_llama_parse:
                somar_pages = 0
            else:
                somar_pages = 1

            page_number = 0
            for start, end, page_number in page_boundaries:
                if start <= start_char < end:
                    page_number = page_number
                    break

            doc_chunk = DocumentChunk(  # Gera o objeto do chunk com informações adicionais, como a posição e id do chunk
                content=chunk,
                contextual_summary="",
                page_number=page_number + somar_pages,  # 1-based page numbering
                chunk_id=chunk_id,
                start_char=char_count + start_char,
                end_char=char_count + end_char,
            )
            chunks.append(doc_chunk)

            # Store metadata for later retrieval
            self.chunk_metadata[chunk_id] = {
                "page": doc_chunk.page_number,
                "start_char": doc_chunk.start_char,
                "end_char": doc_chunk.end_char,
            }

            # char_count += len(text)
        print("TERMINOU DE ORGANIZAR PDFS EM CHUNKS")

        return chunks, initial_chunks

    def load_and_split_text(self, text: str) -> List[DocumentChunk]:
        """Load Text and split into chunks with metadata - Criei essa função apenas para o ragas"""
        page = Document(page_content=text, metadata={"page": 1})
        chunks = []
        char_count = 0

        text = page.page_content
        page_chunks = self.text_splitter.split_text(
            text
        )  # Quebra o item que é um Document de UMA PÁGINA inteira em um lista onde cada item é referente a um chunk, que são pedaços menores do que uma página.
        print("\n\n\npage_chunks: ", page_chunks)

        for chunk in page_chunks:
            chunk_id = str(uuid.uuid4())
            start_char = text.find(
                chunk
            )  # Retorna a posição onde se encontra o chunk dentro da página inteira
            end_char = start_char + len(chunk)

            doc_chunk = DocumentChunk(  # Gera o objeto do chunk com informações adicionais, como a posição e id do chunk
                content=chunk,
                page_number=cast(int, page.metadata.get("page"))
                + 1,  # 1-based page numbering
                chunk_id=chunk_id,
                start_char=char_count + start_char,
                end_char=char_count + end_char,
            )
            chunks.append(doc_chunk)

            # Store metadata for later retrieval
            self.chunk_metadata[chunk_id] = {
                "page": doc_chunk.page_number,
                "start_char": doc_chunk.start_char,
                "end_char": doc_chunk.end_char,
            }

        char_count += len(text)

        return chunks