Blob manager improvements/tests

Author: pamelafox

app/backend/app.py

@@ -561,7 +561,7 @@ async def setup_clients():
         openai_organization=OPENAI_ORGANIZATION,
     )
 
-    user_image_blob_manager = None
+    user_blob_manager = None
     if USE_USER_UPLOAD:
         current_app.logger.info("USE_USER_UPLOAD is true, setting up user upload feature")
         if not AZURE_USERSTORAGE_ACCOUNT or not AZURE_USERSTORAGE_CONTAINER:
@@ -616,7 +616,7 @@ async def setup_clients():
             embeddings=text_embeddings_service,
             image_embeddings=image_embeddings_service,
             search_field_name_embedding=AZURE_SEARCH_FIELD_NAME_EMBEDDING,
-            blob_manager=user_image_blob_manager,
+            blob_manager=user_blob_manager,
         )
         current_app.config[CONFIG_INGESTER] = ingester
 

app/backend/approaches/approach.py

@@ -418,8 +418,8 @@ async def download_blob_as_base64(self, blob_url: str, user_oid: Optional[str] =
         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)
-        blob = await blob_downloader.readall()
-        if blob is not None:
+        if blob_downloader is not None:
+            blob = await blob_downloader.readall()
             img = base64.b64encode(blob).decode("utf-8")
             return f"data:image/png;base64,{img}"
         return None

app/backend/prepdocslib/blobmanager.py

@@ -179,7 +179,7 @@ async def upload_blob(self, file: Union[File, IO], filename: str, user_oid: str)
             str: The URL of the uploaded file, with forward slashes (not URL-encoded)
         """
         # Ensure user directory exists but don't create a subdirectory
-        user_directory_client = await self._ensure_directory(user_oid, owner=user_oid)
+        user_directory_client = await self._ensure_directory(directory_path=user_oid, user_oid=user_oid)
 
         # Create file directly in user directory
         file_client = user_directory_client.get_file_client(filename)
@@ -240,12 +240,12 @@ async def upload_document_image(
         Returns:
             str: The URL of the uploaded file, with forward slashes (not URL-encoded)
         """
-        await self._ensure_directory(user_oid, owner=user_oid)
-        directory_path = self._get_image_directory_path(document_filename, user_oid, image_page_num)
-        image_directory_client = await self._ensure_directory(directory_path, owner=user_oid)
+        await self._ensure_directory(directory_path=user_oid, user_oid=user_oid)
+        image_directory_path = self._get_image_directory_path(document_filename, user_oid, image_page_num)
+        image_directory_client = await self._ensure_directory(directory_path=image_directory_path, user_oid=user_oid)
         file_client = image_directory_client.get_file_client(image_filename)
         image_bytes = BaseBlobManager.add_image_citation(image_bytes, document_filename, image_filename, image_page_num)
-        logger.info("Uploading document image '%s' to '%s'", image_filename, directory_path)
+        logger.info("Uploading document image '%s' to '%s'", image_filename, image_directory_path)
         await file_client.upload_data(image_bytes, overwrite=True, metadata={"UploadedBy": user_oid})
         return unquote(file_client.url)
 
@@ -284,7 +284,7 @@ async def download_blob(
             filename = path_parts[-1]
 
         try:
-            user_directory_client = self._ensure_directory(directory_path, user_oid)
+            user_directory_client = 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
@@ -310,15 +310,17 @@ async def remove_blob(self, filename: str, user_oid: str) -> None:
             ResourceNotFoundError: If the file does not exist
         """
         # Ensure the user directory exists
-        user_directory_client = await self._ensure_directory(user_oid, owner=user_oid)
+        user_directory_client = await self._ensure_directory(directory_path=user_oid, user_oid=user_oid)
         # Delete the main document file
         file_client = user_directory_client.get_file_client(filename)
         await file_client.delete_file()
 
         # Try to delete any associated image directories
         try:
             image_directory_path = self._get_image_directory_path(filename, user_oid)
-            image_directory_client = await self._ensure_directory(image_directory_path, owner=user_oid)
+            image_directory_client = await self._ensure_directory(
+                directory_path=image_directory_path, user_oid=user_oid
+            )
             await image_directory_client.delete_directory()
             logger.info(f"Deleted associated image directory: {image_directory_path}")
         except ResourceNotFoundError:
@@ -338,7 +340,7 @@ async def list_blobs(self, user_oid: str) -> list[str]:
         Returns:
             list[str]: List of filenames that belong to the user
         """
-        user_directory_client = await self._ensure_directory(user_oid, owner=user_oid)
+        user_directory_client = await self._ensure_directory(directory_path=user_oid, user_oid=user_oid)
         files = []
         try:
             all_paths = user_directory_client.get_paths()
@@ -389,21 +391,30 @@ def __init__(
         self.subscription_id = subscription_id
         self.image_container = image_container
 
-    async def upload_blob(self, file: File) -> Optional[list[str]]:
-        async with BlobServiceClient(
+    def get_managedidentity_connectionstring(self):
+        if not self.account or not self.resource_group or not self.subscription_id:
+            raise ValueError("Account, resource group, and subscription ID must be set to generate connection string.")
+        return f"ResourceId=/subscriptions/{self.subscription_id}/resourceGroups/{self.resource_group}/providers/Microsoft.Storage/storageAccounts/{self.account};"
+
+    async def _get_service_client(self) -> BlobServiceClient:
+        return BlobServiceClient(
             account_url=self.endpoint, credential=self.credential, max_single_put_size=4 * 1024 * 1024
-        ) as service_client, service_client.get_container_client(self.container) as container_client:
-            if not await container_client.exists():
-                await container_client.create_container()
-
-            # Re-open and upload the original file
-            if file.url is None:
-                with open(file.content.name, "rb") as reopened_file:
-                    blob_name = self.blob_name_from_file_name(file.content.name)
-                    logger.info("Uploading blob for document '%s'", blob_name)
-                    blob_client = await container_client.upload_blob(blob_name, reopened_file, overwrite=True)
-                    file.url = blob_client.url
-        return None
+        )
+
+    async def upload_blob(self, file: File) -> Optional[list[str]]:
+        blob_service_client = await self._get_service_client()
+        container_client = blob_service_client.get_container_client(self.container)
+        if not await container_client.exists():
+            await container_client.create_container()
+
+        # Re-open and upload the original file
+        if file.url is None:
+            with open(file.content.name, "rb") as reopened_file:
+                blob_name = self.blob_name_from_file_name(file.content.name)
+                logger.info("Uploading blob for document '%s'", blob_name)
+                blob_client = await container_client.upload_blob(blob_name, reopened_file, overwrite=True)
+                file.url = blob_client.url
+        blob_service_client.close()
 
     async def upload_document_image(
         self,
@@ -421,22 +432,16 @@ async def upload_document_image(
             raise ValueError(
                 "user_oid is not supported for BlobManager. Use AdlsBlobManager for user-specific operations."
             )
-        async with BlobServiceClient(
-            account_url=self.endpoint, credential=self.credential, max_single_put_size=4 * 1024 * 1024
-        ) as service_client, service_client.get_container_client(self.image_container) as container_client:
-            if not await container_client.exists():
-                await container_client.create_container()
-            image_bytes = self.add_image_citation(image_bytes, document_filename, image_filename, image_page_num)
-            blob_name = f"{self.blob_name_from_file_name(document_filename)}/page{image_page_num}/{image_filename}"
-            logger.info("Uploading blob for document image '%s'", blob_name)
-            blob_client = await container_client.upload_blob(blob_name, image_bytes, overwrite=True)
-            return blob_client.url
-        return None
-
-    def get_managedidentity_connectionstring(self):
-        if not self.account or not self.resource_group or not self.subscription_id:
-            raise ValueError("Account, resource group, and subscription ID must be set to generate connection string.")
-        return f"ResourceId=/subscriptions/{self.subscription_id}/resourceGroups/{self.resource_group}/providers/Microsoft.Storage/storageAccounts/{self.account};"
+        blob_service_client = await self._get_service_client()
+        container_client = blob_service_client.get_container_client(self.container)
+        if not await container_client.exists():
+            await container_client.create_container()
+        image_bytes = self.add_image_citation(image_bytes, document_filename, image_filename, image_page_num)
+        blob_name = f"{self.blob_name_from_file_name(document_filename)}/page{image_page_num}/{image_filename}"
+        logger.info("Uploading blob for document image '%s'", blob_name)
+        blob_client = await container_client.upload_blob(blob_name, image_bytes, overwrite=True)
+        blob_service_client.close()
+        return blob_client.url
 
     async def download_blob(
         self, blob_path: str, user_oid: Optional[str] = None
@@ -445,43 +450,43 @@ async def download_blob(
             raise ValueError(
                 "user_oid is not supported for BlobManager. Use AdlsBlobManager for user-specific operations."
             )
-
-        async with BlobServiceClient(
-            account_url=self.endpoint, credential=self.credential, max_single_put_size=4 * 1024 * 1024
-        ) as service_client, service_client.get_container_client(self.container) as container_client:
-            if not await container_client.exists():
-                return None
-            blob_client = container_client.get_blob_client(blob_path)
-            try:
-                blob = await blob_client.download_blob()
-                if not blob.properties:
-                    logger.warning(f"No blob exists for {blob_path}")
-                    return None
-                return blob
-            except ResourceNotFoundError:
-                logger.warning("Blob not found: %s", blob_path)
+        blob_service_client = await self._get_service_client()
+        container_client = blob_service_client.get_container_client(self.container)
+        if not await container_client.exists():
+            return None
+        if len(blob_path) == 0:
+            logger.warning("Blob path is empty")
+            return None
+        blob_client = container_client.get_blob_client(blob_path)
+        try:
+            blob = await blob_client.download_blob()
+            if not blob.properties:
+                logger.warning(f"No blob exists for {blob_path}")
                 return None
+            return blob
+        except ResourceNotFoundError:
+            logger.warning("Blob not found: %s", blob_path)
+            return None
+        # TODO: close client properly
 
     async def remove_blob(self, path: Optional[str] = None):
-        async with BlobServiceClient(
-            account_url=self.endpoint, credential=self.credential
-        ) as service_client, service_client.get_container_client(self.container) as container_client:
-            if not await container_client.exists():
-                return
-            if path is None:
-                prefix = None
-                blobs = container_client.list_blob_names()
-            else:
-                prefix = os.path.splitext(os.path.basename(path))[0]
-                blobs = container_client.list_blob_names(name_starts_with=os.path.splitext(os.path.basename(prefix))[0])
-            async for blob_path in blobs:
-                # This still supports PDFs split into individual pages, but we could remove in future to simplify code
-                if (
-                    prefix is not None
-                    and (
-                        not re.match(rf"{prefix}-\d+\.pdf", blob_path) or not re.match(rf"{prefix}-\d+\.png", blob_path)
-                    )
-                ) or (path is not None and blob_path == os.path.basename(path)):
-                    continue
-                logger.info("Removing blob %s", blob_path)
-                await container_client.delete_blob(blob_path)
+        blob_service_client = await self._get_service_client()
+        container_client = blob_service_client.get_container_client(self.container)
+        if not await container_client.exists():
+            return
+        if path is None:
+            prefix = None
+            blobs = container_client.list_blob_names()
+        else:
+            prefix = os.path.splitext(os.path.basename(path))[0]
+            blobs = container_client.list_blob_names(name_starts_with=os.path.splitext(os.path.basename(prefix))[0])
+        async for blob_path in blobs:
+            # This still supports PDFs split into individual pages, but we could remove in future to simplify code
+            if (
+                prefix is not None
+                and (not re.match(rf"{prefix}-\d+\.pdf", blob_path) or not re.match(rf"{prefix}-\d+\.png", blob_path))
+            ) or (path is not None and blob_path == os.path.basename(path)):
+                continue
+            logger.info("Removing blob %s", blob_path)
+            await container_client.delete_blob(blob_path)
+        blob_service_client.close()

tests/conftest.py

@@ -29,9 +29,14 @@
 
 import app
 import core
+from approaches.chatreadretrieveread import ChatReadRetrieveReadApproach
+from approaches.promptmanager import PromptyManager
 from core.authentication import AuthenticationHelper
+from prepdocslib.blobmanager import AdlsBlobManager, BlobManager
 
 from .mocks import (
+    MOCK_EMBEDDING_DIMENSIONS,
+    MOCK_EMBEDDING_MODEL_NAME,
     MockAsyncPageIterator,
     MockAsyncSearchResultsIterator,
     MockAzureCredential,
@@ -1058,3 +1063,37 @@ def mock_get_bearer_token(*args, **kwargs):
         return mock_token_provider
 
     monkeypatch.setattr("azure.identity.aio.get_bearer_token_provider", mock_get_bearer_token)
+
+
+@pytest.fixture
+def chat_approach():
+    return ChatReadRetrieveReadApproach(
+        search_client=None,
+        search_index_name=None,
+        agent_model=None,
+        agent_deployment=None,
+        agent_client=None,
+        auth_helper=None,
+        openai_client=None,
+        chatgpt_model="gpt-4.1-mini",
+        chatgpt_deployment="chat",
+        embedding_deployment="embeddings",
+        embedding_model=MOCK_EMBEDDING_MODEL_NAME,
+        embedding_dimensions=MOCK_EMBEDDING_DIMENSIONS,
+        embedding_field="embedding3",
+        sourcepage_field="",
+        content_field="",
+        query_language="en-us",
+        query_speller="lexicon",
+        prompt_manager=PromptyManager(),
+        user_blob_manager=AdlsBlobManager(
+            endpoint="https://test-userstorage-account.dfs.core.windows.net",
+            container="test-userstorage-container",
+            credential=MockAzureCredential(),
+        ),
+        global_blob_manager=BlobManager(  # on normal Azure storage
+            endpoint="https://test-globalstorage-account.blob.core.windows.net",
+            container="test-globalstorage-container",
+            credential=MockAzureCredential(),
+        ),
+    )

tests/snapshots/test_app/test_chat_vision/client0/result.json

@@ -8,7 +8,7 @@
                 "Financial Market Analysis Report 2023.pdf#page=2"
             ],
             "images": [
-                "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg=="
+                "data:image/png;base64,base64encodedimage"
             ],
             "text": [
                 "Financial Market Analysis Report 2023.pdf#page=7:  This section examines the correlations between stock indices, cryptocurrency prices, and commodity prices, revealing how changes in one market can have ripple effects across the financial ecosystem.### Impact of Macroeconomic Factors   <figure><figcaption>Impact of Interest Rates, Inflation, and GDP Growth on Financial Markets<br>The image is a line graph titled \"on Financial Markets\" displaying data from 2018 to 2023. It tracks three variables: Interest Rates %, Inflation Data %, and GDP Growth %, each represented by a different colored line (blue for Interest Rates, orange for Inflation Data, and gray for GDP Growth). Interest Rates % start around 2% in 2018, dip to about 0.25% in 2021, then rise to 1.5% in 2023. Inflation Data % begin at approximately 1.9% in 2018, rise to a peak near 3.4% in 2022, and then decrease to 2.5% in 2023. GDP Growth % shows significant fluctuations, starting at 3% in 2018, plunging to almost -4% in 2020, then rebounding to around 4.5% in 2021 before gradually declining to around 2.8% in 2023.</figcaption></figure>   Macroeconomic factors such as interest rates, inflation, and GDP growth play a pivotal role in shaping financial markets.",
@@ -156,7 +156,7 @@
                             },
                             {
                                 "image_url": {
-                                    "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg=="
+                                    "url": "data:image/png;base64,base64encodedimage"
                                 },
                                 "type": "image_url"
                             },

tests/test_app.py

@@ -1027,7 +1027,15 @@ async def test_chat_stream_followup(client, snapshot):
 
 
 @pytest.mark.asyncio
-async def test_chat_vision(vision_client, snapshot):
+async def test_chat_vision(monkeypatch, vision_client, snapshot):
+    async def mock_download_blob_as_base64(self, blob_url, user_oid):
+        return "data:image/png;base64,base64encodedimage"
+
+    monkeypatch.setattr(
+        "approaches.chatreadretrieveread.ChatReadRetrieveReadApproach.download_blob_as_base64",
+        mock_download_blob_as_base64,
+    )
+
     response = await vision_client.post(
         "/chat",
         json={"messages": [{"content": "Are interest rates high?", "role": "user"}]},

tests/test_chatapproach.py

@@ -25,30 +25,6 @@ async def mock_retrieval(*args, **kwargs):
     return mock_retrieval_response()
 
 
-@pytest.fixture
-def chat_approach():
-    return ChatReadRetrieveReadApproach(
-        search_client=None,
-        search_index_name=None,
-        agent_model=None,
-        agent_deployment=None,
-        agent_client=None,
-        auth_helper=None,
-        openai_client=None,
-        chatgpt_model="gpt-4.1-mini",
-        chatgpt_deployment="chat",
-        embedding_deployment="embeddings",
-        embedding_model=MOCK_EMBEDDING_MODEL_NAME,
-        embedding_dimensions=MOCK_EMBEDDING_DIMENSIONS,
-        embedding_field="embedding3",
-        sourcepage_field="",
-        content_field="",
-        query_language="en-us",
-        query_speller="lexicon",
-        prompt_manager=PromptyManager(),
-    )
-
-
 def test_get_search_query(chat_approach):
     payload = """
     {