nltk security issue and upgrade unstructured (#9558)

api/core/rag/extractor/unstructured/unstructured_eml_extractor.py

@@ -10,24 +10,26 @@ logger = logging.getLogger(__name__)
 
 
 class UnstructuredEmailExtractor(BaseExtractor):
-    """Load msg files.
+    """Load eml files.
     Args:
         file_path: Path to the file to load.
     """
 
-    def __init__(
-        self,
-        file_path: str,
-        api_url: str,
-    ):
+    def __init__(self, file_path: str, api_url: str, api_key: str):
         """Initialize with file path."""
         self._file_path = file_path
         self._api_url = api_url
+        self._api_key = api_key
 
     def extract(self) -> list[Document]:
-        from unstructured.partition.email import partition_email
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
 
-        elements = partition_email(filename=self._file_path)
+            elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+        else:
+            from unstructured.partition.email import partition_email
+
+            elements = partition_email(filename=self._file_path)
 
         # noinspection PyBroadException
         try:

api/core/rag/extractor/unstructured/unstructured_epub_extractor.py

@@ -19,15 +19,23 @@ class UnstructuredEpubExtractor(BaseExtractor):
         self,
         file_path: str,
         api_url: Optional[str] = None,
+        api_key: Optional[str] = None,
     ):
         """Initialize with file path."""
         self._file_path = file_path
         self._api_url = api_url
+        self._api_key = api_key
 
     def extract(self) -> list[Document]:
-        from unstructured.partition.epub import partition_epub
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
+
+            elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+        else:
+            from unstructured.partition.epub import partition_epub
+
+            elements = partition_epub(filename=self._file_path, xml_keep_tags=True)
 
-        elements = partition_epub(filename=self._file_path, xml_keep_tags=True)
         from unstructured.chunking.title import chunk_by_title
 
         chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000)

api/core/rag/extractor/unstructured/unstructured_markdown_extractor.py

@@ -24,19 +24,21 @@ class UnstructuredMarkdownExtractor(BaseExtractor):
             if the specified encoding fails.
     """
 
-    def __init__(
-        self,
-        file_path: str,
-        api_url: str,
-    ):
+    def __init__(self, file_path: str, api_url: str, api_key: str):
         """Initialize with file path."""
         self._file_path = file_path
         self._api_url = api_url
+        self._api_key = api_key
 
     def extract(self) -> list[Document]:
-        from unstructured.partition.md import partition_md
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
 
-        elements = partition_md(filename=self._file_path)
+            elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+        else:
+            from unstructured.partition.md import partition_md
+
+            elements = partition_md(filename=self._file_path)
         from unstructured.chunking.title import chunk_by_title
 
         chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000)

api/core/rag/extractor/unstructured/unstructured_msg_extractor.py

@@ -14,15 +14,21 @@ class UnstructuredMsgExtractor(BaseExtractor):
         file_path: Path to the file to load.
     """
 
-    def __init__(self, file_path: str, api_url: str):
+    def __init__(self, file_path: str, api_url: str, api_key: str):
         """Initialize with file path."""
         self._file_path = file_path
         self._api_url = api_url
+        self._api_key = api_key
 
     def extract(self) -> list[Document]:
-        from unstructured.partition.msg import partition_msg
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
 
-        elements = partition_msg(filename=self._file_path)
+            elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+        else:
+            from unstructured.partition.msg import partition_msg
+
+            elements = partition_msg(filename=self._file_path)
         from unstructured.chunking.title import chunk_by_title
 
         chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000)

api/core/rag/extractor/unstructured/unstructured_pdf_extractor.py

@@ -0,0 +1,47 @@
+import logging
+
+from core.rag.extractor.extractor_base import BaseExtractor
+from core.rag.models.document import Document
+
+logger = logging.getLogger(__name__)
+
+
+class UnstructuredPDFExtractor(BaseExtractor):
+    """Load pdf files.
+
+
+    Args:
+        file_path: Path to the file to load.
+
+        api_url: Unstructured API URL
+
+        api_key: Unstructured API Key
+    """
+
+    def __init__(self, file_path: str, api_url: str, api_key: str):
+        """Initialize with file path."""
+        self._file_path = file_path
+        self._api_url = api_url
+        self._api_key = api_key
+
+    def extract(self) -> list[Document]:
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
+
+            elements = partition_via_api(
+                filename=self._file_path, api_url=self._api_url, api_key=self._api_key, strategy="auto"
+            )
+        else:
+            from unstructured.partition.pdf import partition_pdf
+
+            elements = partition_pdf(filename=self._file_path, strategy="auto")
+
+        from unstructured.chunking.title import chunk_by_title
+
+        chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000)
+        documents = []
+        for chunk in chunks:
+            text = chunk.text.strip()
+            documents.append(Document(page_content=text))
+
+        return documents

api/core/rag/extractor/unstructured/unstructured_ppt_extractor.py

@@ -7,7 +7,7 @@ logger = logging.getLogger(__name__)
 
 
 class UnstructuredPPTExtractor(BaseExtractor):
-    """Load msg files.
+    """Load ppt files.
 
 
     Args:
@@ -21,9 +21,12 @@ class UnstructuredPPTExtractor(BaseExtractor):
         self._api_key = api_key
 
     def extract(self) -> list[Document]:
-        from unstructured.partition.api import partition_via_api
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
 
-        elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+            elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+        else:
+            raise NotImplementedError("Unstructured API Url is not configured")
         text_by_page = {}
         for element in elements:
             page = element.metadata.page_number

api/core/rag/extractor/unstructured/unstructured_pptx_extractor.py

@@ -7,22 +7,28 @@ logger = logging.getLogger(__name__)
 
 
 class UnstructuredPPTXExtractor(BaseExtractor):
-    """Load msg files.
+    """Load pptx files.
 
 
     Args:
         file_path: Path to the file to load.
     """
 
-    def __init__(self, file_path: str, api_url: str):
+    def __init__(self, file_path: str, api_url: str, api_key: str):
         """Initialize with file path."""
         self._file_path = file_path
         self._api_url = api_url
+        self._api_key = api_key
 
     def extract(self) -> list[Document]:
-        from unstructured.partition.pptx import partition_pptx
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
 
-        elements = partition_pptx(filename=self._file_path)
+            elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+        else:
+            from unstructured.partition.pptx import partition_pptx
+
+            elements = partition_pptx(filename=self._file_path)
         text_by_page = {}
         for element in elements:
             page = element.metadata.page_number

api/core/rag/extractor/unstructured/unstructured_xml_extractor.py

@@ -7,22 +7,29 @@ logger = logging.getLogger(__name__)
 
 
 class UnstructuredXmlExtractor(BaseExtractor):
-    """Load msg files.
+    """Load xml files.
 
 
     Args:
         file_path: Path to the file to load.
     """
 
-    def __init__(self, file_path: str, api_url: str):
+    def __init__(self, file_path: str, api_url: str, api_key: str):
         """Initialize with file path."""
         self._file_path = file_path
         self._api_url = api_url
+        self._api_key = api_key
 
     def extract(self) -> list[Document]:
-        from unstructured.partition.xml import partition_xml
+        if self._api_url:
+            from unstructured.partition.api import partition_via_api
+
+            elements = partition_via_api(filename=self._file_path, api_url=self._api_url, api_key=self._api_key)
+        else:
+            from unstructured.partition.xml import partition_xml
+
+            elements = partition_xml(filename=self._file_path, xml_keep_tags=True)
 
-        elements = partition_xml(filename=self._file_path, xml_keep_tags=True)
         from unstructured.chunking.title import chunk_by_title
 
         chunks = chunk_by_title(elements, max_characters=2000, combine_text_under_n_chars=2000)