Repository: StanGirard/quivr
File: core/quivr_core/processor/implementations/megaparse_processor.py (Python)

Code before:
        self.enc = tiktoken.get_encoding("cl100k_base")
        self.splitter_config = splitter_config
        self.megaparse_config = megaparse_config

        if splitter:
            self.text_splitter = splitter
        else:
            self.text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
                chunk_size=splitter_config.chunk_size,
                chunk_overlap=splitter_config.chunk_overlap,
            )

    @property
    def processor_metadata(self):
        return {
            "chunk_overlap": self.splitter_config.chunk_overlap,
        }

    async def process_file_inner(self, file: QuivrFile) -> list[Document]:
        logger.info(f"Uploading file {file.path} to MegaParse")
        parser = UnstructuredParser(**self.megaparse_config.model_dump())
        megaparse = MegaParse(parser)
        response = await megaparse.aload(file.path)
        logger.info(f"File :  {response}")
        document = Document(
            page_content=response,
        )
        if len(response) > self.splitter_config.chunk_size:
            docs = self.text_splitter.split_documents([document])
            for doc in docs:
                doc.metadata = {"chunk_size": len(self.enc.encode(doc.page_content))}
            return docs
        return [document]

    # async def process_file_inner(self, file: QuivrFile) -> list[Document]:
    #     api_key = str(os.getenv("MEGAPARSE_API_KEY"))
    #     megaparse = MegaParseSDK(api_key)
    #     logger.info(f"Uploading file {file.path} to MegaParse")
    #     data = {
    #         "method": self.megaparse_config.method,
    #         "strategy": self.megaparse_config.strategy,
    #         "check_table": self.megaparse_config.check_table,
    #         "parsing_instruction": self.megaparse_config.parsing_instruction,
    #         "model_name": self.megaparse_config.model_name,
    #     }
    #     response = await megaparse.file.upload(
    #         file_path=str(file.path),
    #         **data,
    #     )
    #     document = Document(
    #         page_content=response["result"],

Review comment (bug):
We still have a bug of not having chunk_size in the metadata if content is smaller. We should always call the splitter the overhead is minimal for small chunks

Code after:
        self.enc = tiktoken.get_encoding("cl100k_base")
        self.splitter_config = splitter_config
        self.megaparse_config = megaparse_config

        if splitter:
            self.text_splitter = splitter
        else:
            self.text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
                chunk_size=splitter_config.chunk_size,
                chunk_overlap=splitter_config.chunk_overlap,
            )

    @property
    def processor_metadata(self):
        return {
            "chunk_overlap": self.splitter_config.chunk_overlap,
        }

    async def process_file_inner(self, file: QuivrFile) -> list[Document]:
        logger.info(f"Uploading file {file.path} to MegaParse")
        parser = UnstructuredParser(**self.megaparse_config.model_dump())
        megaparse = MegaParse(parser)
        response = await megaparse.aload(file.path)
        logger.info(f"File :  {response}")
        document = Document(
            page_content=response,
        )

        docs = self.text_splitter.split_documents([document])
        for doc in docs:
            doc.metadata = {"chunk_size": len(self.enc.encode(doc.page_content))}
        return docs

    # async def process_file_inner(self, file: QuivrFile) -> list[Document]:
    #     api_key = str(os.getenv("MEGAPARSE_API_KEY"))
    #     megaparse = MegaParseSDK(api_key)
    #     logger.info(f"Uploading file {file.path} to MegaParse")
    #     data = {
    #         "method": self.megaparse_config.method,
    #         "strategy": self.megaparse_config.strategy,
    #         "check_table": self.megaparse_config.check_table,
    #         "parsing_instruction": self.megaparse_config.parsing_instruction,
    #         "model_name": self.megaparse_config.model_name,
    #     }
    #     response = await megaparse.file.upload(
    #         file_path=str(file.path),
    #         **data,
    #     )
    #     document = Document(
    #         page_content=response["result"],
    #     )