Add as_bytes option

Author: pamelafox

app/backend/approaches/approach.py

@@ -414,13 +414,13 @@ async def download_blob_as_base64(self, blob_url: str, user_oid: Optional[str] =
             blob_path = blob_url
 
         # Download the blob using the appropriate client
-        blob_manager = self.global_blob_manager
+        blob_bytes = None
         if ".dfs.core.windows.net" in blob_url and self.user_blob_manager:
-            blob_manager = self.user_blob_manager
-        blob_downloader = await blob_manager.download_blob(blob_path, user_oid=user_oid)
-        if blob_downloader is not None:
-            blob = await blob_downloader.readall()
-            img = base64.b64encode(blob).decode("utf-8")
+            blob_bytes = await self.user_blob_manager.download_blob(blob_path, user_oid=user_oid, as_bytes=True)
+        elif self.global_blob_manager:
+            blob_bytes = await self.global_blob_manager.download_blob(blob_path, as_bytes=True)
+        if blob_bytes and isinstance(blob_bytes, (bytes)):
+            img = base64.b64encode(blob_bytes).decode("utf-8")
             return f"data:image/png;base64,{img}"
         return None
 

app/backend/prepdocslib/blobmanager.py

@@ -103,8 +103,8 @@ async def upload_document_image(
         raise NotImplementedError("Subclasses must implement this method")
 
     async def download_blob(
-        self, blob_path: str, user_oid: Optional[str] = None
-    ) -> Optional[Union[BlobStorageStreamDownloader, AdlsBlobStorageStreamDownloader]]:
+        self, blob_path: str, user_oid: Optional[str] = None, as_bytes: bool = False
+    ) -> Optional[Union[BlobStorageStreamDownloader, AdlsBlobStorageStreamDownloader, bytes]]:
         """
         Downloads a blob from Azure Storage.
         If user_oid is provided, it checks if the blob belongs to the user.
@@ -253,7 +253,7 @@ async def upload_document_image(
         return unquote(file_client.url)
 
     async def download_blob(
-        self, blob_path: str, user_oid: Optional[str] = None
+        self, blob_path: str, user_oid: Optional[str] = None, as_bytes: bool = False
     ) -> Optional[Union[AdlsBlobStorageStreamDownloader, BlobStorageStreamDownloader]]:
         """
         Downloads a blob from Azure Data Lake Storage.
@@ -290,7 +290,10 @@ async def download_blob(
             user_directory_client = await self._ensure_directory(directory_path=directory_path, user_oid=user_oid)
             file_client = user_directory_client.get_file_client(filename)
             blob = await file_client.download_file()
-            return blob
+            if as_bytes:
+                return await blob.readall()
+            else:
+                return blob
         except ResourceNotFoundError:
             logger.warning(f"Directory or file not found: {directory_path}/{filename}")
             return None
@@ -444,8 +447,8 @@ async def upload_document_image(
         return blob_client.url
 
     async def download_blob(
-        self, blob_path: str, user_oid: Optional[str] = None
-    ) -> Optional[Union[BlobStorageStreamDownloader, AdlsBlobStorageStreamDownloader]]:
+        self, blob_path: str, user_oid: Optional[str] = None, as_bytes: bool = False
+    ) -> Optional[Union[BlobStorageStreamDownloader, AdlsBlobStorageStreamDownloader, bytes]]:
         if user_oid is not None:
             raise ValueError(
                 "user_oid is not supported for BlobManager. Use AdlsBlobManager for user-specific operations."
@@ -462,7 +465,10 @@ async def download_blob(
             if not blob.properties:
                 logger.warning(f"No blob exists for {blob_path}")
                 return None
-            return blob
+            if as_bytes:
+                return await blob.readall()
+            else:
+                return blob
         except ResourceNotFoundError:
             logger.warning("Blob not found: %s", blob_path)
             return None