feat: adicionado suporte para .odt e .txt

_utils/bubble_integrations/obter_arquivo.py

@@ -1,13 +1,14 @@
 # from setup.easy_imports import PyPDFLoader
 import os
-from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
+from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader
 import tempfile
 import requests
 
 from _utils.handle_files import return_document_list_with_llama_parser
-from _utils.langchain_utils.splitter_util import SplitterUtils
+from _utils.langchain_utils.splitter_util import Splitter_Simple, SplitterUtils
 
 splitter_utils = SplitterUtils()
+splitter_simple = Splitter_Simple()
 headers = {"Authorization": f"Bearer {os.environ.get("BUBBLE_TOKEN")}"}
 
 
@@ -36,20 +37,27 @@ async def get_pdf_from_bubble(
     else:
         extension = file_url.split(".")[-1]
         if extension.lower() == "pdf":
-            result = PyPDFLoader(file_url, headers=headers)
+            result = PyPDFLoader(file_url, headers=headers).load()
+        elif extension.lower() == "odt":
+            temp_path = download_file_from_bubble(file_url, headers, ".odt")
+            full_text = splitter_utils.load_odt_file(temp_path)
+            result = splitter_simple.load_and_split_text(full_text)
+        elif extension.lower() == "txt":
+            temp_path = download_file_from_bubble(file_url, headers, ".txt")
+            result = TextLoader(temp_path).load()
         else:
-            temp_path = download_docx(file_url, headers)
-            result = Docx2txtLoader(temp_path)
+            temp_path = download_file_from_bubble(file_url, headers, ".docx")
+            result = Docx2txtLoader(temp_path).load()
 
-        return result.load()
+        return result
 
 
-def download_docx(url, headers):
+def download_file_from_bubble(url, headers, extension: str):
     response = requests.get(url, headers=headers)
     response.raise_for_status()  # Raise an exception for bad responses (status codes 4xx or 5xx)
 
     # Save the downloaded file into a temporary file
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".docx")
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=extension)
     with open(temp_file.name, "wb") as f:
         f.write(response.content)  # por enquanto este arquivo não está sendo excluído
 

_utils/langchain_utils/Splitter_class.py

@@ -1,6 +1,7 @@
 from _utils.bubble_integrations.obter_arquivo import get_pdf_from_bubble
 from _utils.handle_files import return_document_list_with_llama_parser
 from _utils.langchain_utils.splitter_util import (
+    Splitter_Simple,
     SplitterUtils,
     combine_documents_without_losing_pagination,
 )
@@ -9,6 +10,7 @@ from setup.easy_imports import (
     RecursiveCharacterTextSplitter,
     Document,
     Docx2txtLoader,
+    TextLoader,
 )
 from typing import Any, List, Dict, Tuple, Optional, cast
 from _utils.models.gerar_relatorio import (
@@ -16,8 +18,6 @@ from _utils.models.gerar_relatorio import (
 )
 import uuid
 
-splitter_utils = SplitterUtils()
-
 
 class Splitter:
     def __init__(
@@ -25,6 +25,7 @@ class Splitter:
         chunk_size,
         chunk_overlap,
     ):
+        self.splitter_util = SplitterUtils()
         self.splitter_simple = Splitter_Simple(chunk_size, chunk_overlap)
         self.text_splitter = RecursiveCharacterTextSplitter(
             chunk_size=chunk_size, chunk_overlap=chunk_overlap
@@ -56,11 +57,6 @@ class Splitter:
                     pages
                 )
             )
-            # for page in pages:
-            #     full_text_as_string = full_text_as_string + page.page_content
-            # chunks_of_string_only = chunks_of_string_only + self.text_splitter.split_text(
-            #     combined_text
-            # )
         else:
             if should_use_llama_parse:
                 print("\nENVIANDO PDFS PARA LLAMA PARSE")
@@ -73,10 +69,15 @@ class Splitter:
                 )
             else:
                 print("\nCOMEÇANDO LEITURA DO PDF")
-                file_extension = splitter_utils.get_file_type(pdf_path)
+                file_extension = self.splitter_util.get_file_type(pdf_path)
                 print("file_extension: ", file_extension)
                 if file_extension == "pdf":
                     pages = PyPDFLoader(pdf_path).load()
+                elif file_extension == "odt":
+                    full_text = self.splitter_util.load_odt_file(pdf_path)
+                    pages = self.splitter_simple.load_and_split_text(full_text)
+                elif file_extension == "txt":
+                    pages = TextLoader(pdf_path).load()
                 else:
                     pages = Docx2txtLoader(pdf_path).load()
                 print("TERMINOU LEITURA DO PDF")
@@ -177,40 +178,3 @@ class Splitter:
         char_count += len(text)
 
         return chunks
-
-
-class Splitter_Simple:
-    def __init__(
-        self,
-        chunk_size=1000,
-        chunk_overlap=400,
-    ):
-        self.text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=chunk_size, chunk_overlap=chunk_overlap
-        )
-
-    async def load_and_split_document(self, pdf_path: str):
-        """Load PDF and split into chunks with metadata"""
-        print("\nCOMEÇANDO LEITURA DO PDF")
-        pages = PyPDFLoader(pdf_path).load_and_split(self.text_splitter)
-        print("\nTERMINADO LEITURA DO PDF")
-
-        return pages
-
-    def load_and_split_text(self, text: str) -> List[Document]:
-        documents: List[Document] = []
-        chunks = self.text_splitter.split_text(text)
-
-        for chunk in chunks:
-            documents.append(Document(page_content=chunk))
-
-        return documents
-
-    def get_chunks_of_string_only_from_list_of_documents(
-        self, lista_de_documentos: List[Document]
-    ):
-        full_text_as_string = ""
-        for page in lista_de_documentos:
-            full_text_as_string = full_text_as_string + page.page_content
-        full_text_as_array = self.text_splitter.split_text(full_text_as_string)
-        return full_text_as_array

_utils/langchain_utils/splitter_util.py

@@ -1,6 +1,13 @@
 import os
 from typing import List, Tuple
 from langchain_core.documents import Document
+from odf.opendocument import load
+from odf.text import P
+from typing import List
+from setup.easy_imports import (
+    PyPDFLoader,
+    RecursiveCharacterTextSplitter,
+)
 
 
 class SplitterUtils:
@@ -11,10 +18,57 @@ class SplitterUtils:
             return "pdf"
         elif ext == ".docx":
             return "word"
+        elif ext == ".odt":
+            return "odt"
+        elif ext == ".txt":
+            return "txt"
         else:
             print("\next", ext)
             return "unknown"
 
+    def load_odt_file(self, file_path: str):
+        textdoc = load(file_path)
+        all_paragraphs = textdoc.getElementsByType(P)
+        text = "\n".join([p.firstChild.data for p in all_paragraphs if p.firstChild])
+        return text
+
+
+class Splitter_Simple:
+    def __init__(
+        self,
+        chunk_size=1000,
+        chunk_overlap=400,
+    ):
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size, chunk_overlap=chunk_overlap
+        )
+
+    async def load_and_split_document(self, pdf_path: str):
+        """Load PDF and split into chunks with metadata"""
+        print("\nCOMEÇANDO LEITURA DO PDF")
+        pages = PyPDFLoader(pdf_path).load_and_split(self.text_splitter)
+        print("\nTERMINADO LEITURA DO PDF")
+
+        return pages
+
+    def load_and_split_text(self, text: str) -> List[Document]:
+        documents: List[Document] = []
+        chunks = self.text_splitter.split_text(text)
+
+        for chunk in chunks:
+            documents.append(Document(page_content=chunk))
+
+        return documents
+
+    def get_chunks_of_string_only_from_list_of_documents(
+        self, lista_de_documentos: List[Document]
+    ):
+        full_text_as_string = ""
+        for page in lista_de_documentos:
+            full_text_as_string = full_text_as_string + page.page_content
+        full_text_as_array = self.text_splitter.split_text(full_text_as_string)
+        return full_text_as_array
+
 
 def combine_documents_without_losing_pagination(documents: list[Document]):
     combined_text = ""

setup/easy_imports.py

@@ -13,7 +13,7 @@ from langchain_huggingface import HuggingFaceEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain_core.runnables import RunnablePassthrough
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
+from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader
 from langchain_community.vectorstores import Chroma
 from langchain_google_genai import ChatGoogleGenerativeAI