Add support for archive files

Author: vegito22

llmstack/contrib/data/pipelines/archive.yml

@@ -0,0 +1,16 @@
+name: Archive
+slug: archive
+description: Read data from an archive
+pipeline:
+  source:
+    slug: archive
+    provider_slug: promptly
+  transformations:
+    - slug: splitter
+      provider_slug: unstructured
+  embedding:
+    slug: embeddings-generator
+    provider_slug: promptly
+  destination:
+    slug: vector-store
+    provider_slug: promptly

llmstack/data/apis.py

@@ -75,7 +75,12 @@ def templates(self, request):
         return DRFResponse(templates_data)
 
     def sources(self, request):
-        from llmstack.data.sources import FileSchema, TextSchema, URLSchema
+        from llmstack.data.sources import (
+            ArchiveFileSchema,
+            FileSchema,
+            TextSchema,
+            URLSchema,
+        )
 
         return DRFResponse(
             [
@@ -97,6 +102,12 @@ def sources(self, request):
                     "schema": URLSchema.get_schema(),
                     "ui_schema": URLSchema.get_ui_schema(),
                 },
+                {
+                    "slug": ArchiveFileSchema.slug(),
+                    "provider_slug": ArchiveFileSchema.provider_slug(),
+                    "schema": ArchiveFileSchema.get_schema(),
+                    "ui_schema": ArchiveFileSchema.get_ui_schema(),
+                },
             ]
         )
 

llmstack/data/sources/__init__.py

@@ -1,5 +1,6 @@
 from functools import cache
 
+from llmstack.data.sources.files.archive import ArchiveFileSchema
 from llmstack.data.sources.files.csv import CSVFileSchema
 from llmstack.data.sources.files.file import FileSchema
 from llmstack.data.sources.files.pdf import PdfSchema
@@ -11,7 +12,7 @@
 
 @cache
 def get_source_cls(slug, provider_slug):
-    for cls in [CSVFileSchema, FileSchema, PdfSchema, URLSchema, TextSchema]:
+    for cls in [CSVFileSchema, FileSchema, PdfSchema, URLSchema, TextSchema, ArchiveFileSchema]:
         if cls.slug() == slug and cls.provider_slug() == provider_slug:
             return cls
 

llmstack/data/sources/files/archive.py

@@ -0,0 +1,143 @@
+import base64
+import io
+import logging
+import mimetypes
+import tarfile
+import uuid
+import zipfile
+from typing import List
+
+from pydantic import Field
+
+from llmstack.common.utils.text_extract import extract_text_elements
+from llmstack.common.utils.utils import validate_parse_data_uri
+from llmstack.data.sources.base import BaseSource, DataDocument
+from llmstack.data.sources.utils import (
+    create_source_document_asset,
+    get_source_document_asset_by_objref,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def extract_archive_files(mime_type, file_name, file_data):
+    extracted_files = []
+    if mime_type == "application/zip":
+        with zipfile.ZipFile(io.BytesIO(base64.b64decode(file_data))) as archive:
+            for file_info in archive.infolist():
+                if file_info.is_dir() or file_info.file_size == 0 or file_info.filename.startswith("__MACOSX"):
+                    continue
+                with archive.open(file_info) as file:
+                    file_mime_type = mimetypes.guess_type(file_info.filename)[0]
+                    data_uri = f"data:{file_mime_type};name={file_info.filename};base64,{base64.b64encode(file.read()).decode()}"
+                    extracted_files.append(data_uri)
+    elif mime_type in ["application/x-tar", "application/gzip", "application/x-bzip2"]:
+        with tarfile.open(fileobj=io.BytesIO(file_data), mode="r:*") as archive:
+            for member in archive.getmembers():
+                if member.isfile():
+                    file = archive.extractfile(member)
+                    file_mime_type = mimetypes.guess_type(member.name)[0]
+                    data_uri = (
+                        f"data:{file_mime_type};name={member.name};base64,{base64.b64encode(file.read()).decode()}"
+                    )
+                    extracted_files.append(data_uri)
+    else:
+        logger.warning(f"Unsupported archive mime type: {mime_type}")
+    return extracted_files
+
+
+class ArchiveFileSchema(BaseSource):
+    file: str = Field(
+        description="File to be processed",
+        json_schema_extra={
+            "advanced_parameter": False,
+            "widget": "file",
+            "maxSize": 25000000,
+            "maxFiles": 1,
+            "accepts": {
+                "application/zip": [],
+            },
+        },
+    )
+    split_files: bool = Field(
+        default=False,
+        description="Split the archive into individual files",
+        json_schema_extra={"advanced_parameter": True},
+    )
+
+    @classmethod
+    def slug(cls):
+        return "archive"
+
+    @classmethod
+    def provider_slug(cls):
+        return "promptly"
+
+    def get_data_documents(self, **kwargs) -> List[DataDocument]:
+        if self.split_files:
+            files = extract_archive_files(*validate_parse_data_uri(self.file))
+        else:
+            files = [self.file]
+
+        documents = []
+        for file in files:
+            file_id = str(uuid.uuid4())
+            mime_type, file_name, file_data = validate_parse_data_uri(file)
+            file_objref = create_source_document_asset(
+                file, datasource_uuid=kwargs["datasource_uuid"], document_id=file_id
+            )
+            documents.append(
+                DataDocument(
+                    id_=file_id,
+                    name=file_name,
+                    content=file_objref,
+                    mimetype=mime_type,
+                    metadata={
+                        "file_name": file_name,
+                        "mime_type": mime_type,
+                        "source": file_name,
+                        "datasource_uuid": kwargs["datasource_uuid"],
+                    },
+                    datasource_uuid=kwargs["datasource_uuid"],
+                    extra_info={"extra_data": self.get_extra_data()},
+                )
+            )
+        return documents
+
+    @classmethod
+    def process_document(cls, document: DataDocument) -> DataDocument:
+        data_uri = get_source_document_asset_by_objref(document.content)
+        mime_type, file_name, file_data = validate_parse_data_uri(data_uri)
+
+        if mime_type == "application/zip":
+            extracted_files = extract_archive_files(mime_type, file_name, file_data)
+            elements = []
+            text_content = ""
+            for extracted_file in extracted_files:
+                mime_type, file_name, extracted_file_data = validate_parse_data_uri(extracted_file)
+                text_content += f"File: {file_name}\n"
+                decoded_file_data = base64.b64decode(extracted_file_data)
+                elements += extract_text_elements(
+                    mime_type=mime_type,
+                    data=decoded_file_data,
+                    file_name=file_name,
+                    extra_params=None,
+                )
+                text_content += "".join([element.text for element in elements])
+            text_content += "\n\n"
+        else:
+            decoded_file_data = base64.b64decode(file_data)
+            elements = extract_text_elements(
+                mime_type=mime_type, data=decoded_file_data, file_name=file_name, extra_params=None
+            )
+            text_content = "".join([element.text for element in elements])
+
+        text_data_uri = (
+            f"data:text/plain;name={document.id_}_text.txt;base64,{base64.b64encode(text_content.encode()).decode()}"
+        )
+        text_file_objref = create_source_document_asset(
+            text_data_uri,
+            datasource_uuid=document.metadata["datasource_uuid"],
+            document_id=str(uuid.uuid4()),
+        )
+        return document.model_copy(update={"text": text_content, "text_objref": text_file_objref})