Fix objref lookup in data sources

Author: vegito22

llmstack/data/apis.py

@@ -574,12 +574,14 @@ def delete(self, request, uid):
         datasource.delete()
         return DRFResponse(status=204)
 
-    def process_add_entry_request(self, datasource: DataSource, source_data) -> List[DataDocument]:
+    def process_add_entry_request(self, datasource: DataSource, source_data, request=None) -> List[DataDocument]:
         source_cls = datasource.pipeline_obj.source_cls
         if not source_cls:
             raise Exception("No source class found for data source")
         source = source_cls(**source_data)
-        return source.get_data_documents(datasource_uuid=str(datasource.uuid))
+        return source.get_data_documents(
+            datasource_uuid=str(datasource.uuid), request_user=request.user if request else None
+        )
 
     def add_entry(self, request, uid):
         datasource = get_object_or_404(DataSource, uuid=uuid.UUID(uid))
@@ -594,7 +596,7 @@ def add_entry(self, request, uid):
         if not source_data:
             return DRFResponse({"errors": ["No source_data provided"]}, status=400)
 
-        documents = self.process_add_entry_request(datasource, source_data)
+        documents = self.process_add_entry_request(datasource, source_data, request=request)
         for document in documents:
             create_result = DataSourceEntryViewSet().create_entry(user=request.user, document=document)
             process_result = DataSourceEntryViewSet().process_entry(request=None, uid=str(create_result.data["uuid"]))

llmstack/data/sources/files/archive.py

@@ -86,7 +86,9 @@ 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"])
+            archive_file = get_document_data_uri_from_objref(
+                archive_file, datasource_uuid=kwargs["datasource_uuid"], request_user=kwargs["request_user"]
+            )
 
         if self.split_files:
             files = extract_archive_files(*validate_parse_data_uri(archive_file))

llmstack/data/sources/files/file.py

@@ -55,7 +55,11 @@ def get_data_documents(self, **kwargs) -> List[DataDocument]:
                 list(
                     map(
                         lambda entry: (
-                            get_document_data_uri_from_objref(file_objref, datasource_uuid=kwargs["datasource_uuid"])
+                            get_document_data_uri_from_objref(
+                                file_objref,
+                                datasource_uuid=kwargs["datasource_uuid"],
+                                request_user=kwargs["request_user"],
+                            )
                             if entry.startswith("objref://")
                             else entry
                         ),

llmstack/data/sources/utils.py

@@ -59,36 +59,19 @@ def create_source_document_asset(file, datasource_uuid, document_id):
 
 
 def get_source_document_asset_by_objref(objref):
+    from llmstack.assets.utils import get_asset_by_objref_internal
     from llmstack.data.models import DataSourceEntryFiles
 
-    if not objref:
-        return None
-    asset = None
-    try:
-        category, uuid = objref.strip().split("//")[1].split("/")
-        asset = DataSourceEntryFiles.get_asset_data_uri(DataSourceEntryFiles.objects.get(uuid=uuid), include_name=True)
-    except Exception:
-        pass
-
+    asset_obj = get_asset_by_objref_internal(objref)
+    asset = DataSourceEntryFiles.get_asset_data_uri(asset_obj, include_name=True)
     return asset
 
 
-def get_document_data_uri_from_objref(objref, datasource_uuid):
+def get_document_data_uri_from_objref(objref, datasource_uuid, request_user):
+    from llmstack.assets.utils import get_asset_by_objref
     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
+    asset_obj = get_asset_by_objref(objref, request_user, None)
+    asset = DataSourceEntryFiles.get_asset_data_uri(asset_obj, include_name=True) if asset_obj else None
 
     return asset