Update archive and file source to accept objref

Author: vegito22

llmstack/data/sources/files/archive.py

@@ -14,6 +14,7 @@
 from llmstack.data.sources.base import BaseSource, DataDocument
 from llmstack.data.sources.utils import (
     create_source_document_asset,
+    get_document_data_uri_from_objref,
     get_source_document_asset_by_objref,
 )
 
@@ -74,10 +75,15 @@ def provider_slug(cls):
         return "promptly"
 
     def get_data_documents(self, **kwargs) -> List[DataDocument]:
+        archive_file = self.file
+        # If objref:// is present, get the data URI from the objref
+        if archive_file and archive_file.startswith("objref://"):
+            archive_file = get_document_data_uri_from_objref(archive_file, datasource_uuid=kwargs["datasource_uuid"])
+
         if self.split_files:
-            files = extract_archive_files(*validate_parse_data_uri(self.file))
+            files = extract_archive_files(*validate_parse_data_uri(archive_file))
         else:
-            files = [self.file]
+            files = [archive_file]
 
         documents = []
         for file in files:

llmstack/data/sources/files/file.py

@@ -10,6 +10,7 @@
 from llmstack.data.sources.base import BaseSource, DataDocument
 from llmstack.data.sources.utils import (
     create_source_document_asset,
+    get_document_data_uri_from_objref,
     get_source_document_asset_by_objref,
 )
 
@@ -48,6 +49,21 @@ def provider_slug(cls):
 
     def get_data_documents(self, **kwargs) -> List[DataDocument]:
         files = self.file.split("|")
+        files = list(
+            filter(
+                lambda entry: entry is not None,
+                list(
+                    map(
+                        lambda entry: (
+                            get_document_data_uri_from_objref(file_objref, datasource_uuid=kwargs["datasource_uuid"])
+                            if entry.startswith("objref://")
+                            else entry
+                        ),
+                        files,
+                    )
+                ),
+            )
+        )
         documents = []
         for file in files:
             file_id = str(uuid.uuid4())

llmstack/data/sources/utils.py

@@ -71,3 +71,24 @@ def get_source_document_asset_by_objref(objref):
         pass
 
     return asset
+
+
+def get_document_data_uri_from_objref(objref, datasource_uuid):
+    from llmstack.data.models import DataSourceEntryFiles
+
+    if not objref:
+        return None
+    asset = None
+    try:
+        _, uuid = objref.strip().split("//")[1].split("/")
+        asset_obj = DataSourceEntryFiles.objects.get(uuid=uuid)
+
+        if asset_obj.metadata.get("datasource_uuid") != datasource_uuid:
+            return None
+
+        asset = DataSourceEntryFiles.get_asset_data_uri(asset_obj, include_name=True)
+
+    except Exception:
+        pass
+
+    return asset