More mypy fixes

Author: pamelafox

app/backend/app.py

@@ -25,8 +25,6 @@
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.indexes.aio import SearchIndexClient
-from azure.storage.blob.aio import StorageStreamDownloader as BlobDownloader
-from azure.storage.filedatalake.aio import StorageStreamDownloader as DatalakeDownloader
 from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
 from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
 from opentelemetry.instrumentation.httpx import (
@@ -149,24 +147,33 @@ async def content_file(path: str, auth_claims: dict[str, Any]):
         path = path_parts[0]
     current_app.logger.info("Opening file %s", path)
     blob_manager: BlobManager = current_app.config[CONFIG_GLOBAL_BLOB_MANAGER]
-    blob: Union[BlobDownloader, DatalakeDownloader]
-    blob = await blob_manager.download_blob(path)
-    if blob is None:
+
+    # Get bytes and properties from the blob manager
+    result = await blob_manager.download_blob(path)
+
+    if result is None:
         current_app.logger.info("Path not found in general Blob container: %s", path)
         if current_app.config[CONFIG_USER_UPLOAD_ENABLED]:
             user_oid = auth_claims["oid"]
             user_blob_manager: AdlsBlobManager = current_app.config[CONFIG_USER_BLOB_MANAGER]
-            blob = await user_blob_manager.download_blob(path, user_oid=user_oid)
-            if blob is None:
+            result = await user_blob_manager.download_blob(path, user_oid=user_oid)
+            if result is None:
                 current_app.logger.exception("Path not found in DataLake: %s", path)
-    if not blob or not blob.properties or not blob.properties.has_key("content_settings"):
+
+    if not result:
+        abort(404)
+
+    content, properties = result
+
+    if not properties or "content_settings" not in properties:
         abort(404)
-    mime_type = blob.properties["content_settings"]["content_type"]
+
+    mime_type = properties["content_settings"]["content_type"]
     if mime_type == "application/octet-stream":
         mime_type = mimetypes.guess_type(path)[0] or "application/octet-stream"
-    blob_file = io.BytesIO()
-    await blob.readinto(blob_file)
-    blob_file.seek(0)
+
+    # Create a BytesIO object from the bytes
+    blob_file = io.BytesIO(content)
     return await send_file(blob_file, mimetype=mime_type, as_attachment=False, attachment_filename=path)
 
 

app/backend/approaches/approach.py

@@ -392,7 +392,6 @@ async def download_blob_as_base64(self, blob_url: str, user_oid: Optional[str] =
         Downloads a blob from either Azure Blob Storage or Azure Data Lake Storage and returns it as a base64 encoded string.
 
         Args:
-            storage_client: Either a ContainerClient (for Blob Storage) or FileSystemClient (for Data Lake Storage)
             blob_url: The URL or path to the blob to download
             user_oid: The user's object ID, required for Data Lake Storage operations and access control
 
@@ -414,13 +413,15 @@ 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_bytes = None
+        result = None
         if ".dfs.core.windows.net" in blob_url and self.user_blob_manager:
-            blob_bytes = await self.user_blob_manager.download_blob(blob_path, user_oid=user_oid, as_bytes=True)
+            result = await self.user_blob_manager.download_blob(blob_path, user_oid=user_oid)
         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")
+            result = await self.global_blob_manager.download_blob(blob_path)
+
+        if result:
+            content, _ = result  # Unpack the tuple, ignoring properties
+            img = base64.b64encode(content).decode("utf-8")
             return f"data:image/png;base64,{img}"
         return None
 

app/backend/prepdocslib/blobmanager.py

@@ -3,18 +3,15 @@
 import os
 import re
 from pathlib import Path
-from typing import IO, Optional, Union
+from typing import IO, Any, Optional, TypedDict, Union
 from urllib.parse import unquote
 
 from azure.core.credentials_async import AsyncTokenCredential
 from azure.core.exceptions import ResourceNotFoundError
 from azure.storage.blob.aio import BlobServiceClient
-from azure.storage.blob.aio import (
-    StorageStreamDownloader as BlobStorageStreamDownloader,
-)
-from azure.storage.filedatalake.aio import DataLakeDirectoryClient, FileSystemClient
 from azure.storage.filedatalake.aio import (
-    StorageStreamDownloader as AdlsBlobStorageStreamDownloader,
+    DataLakeDirectoryClient,
+    FileSystemClient,
 )
 from PIL import Image, ImageDraw, ImageFont
 
@@ -23,6 +20,12 @@
 logger = logging.getLogger("scripts")
 
 
+class BlobProperties(TypedDict, total=False):
+    """Properties of a blob, with optional fields for content settings"""
+
+    content_settings: dict[str, Any]
+
+
 class BaseBlobManager:
     """
     Base class for Azure Storage operations, providing common file naming and path utilities
@@ -97,13 +100,13 @@ async def upload_document_image(
         image_bytes: bytes,
         image_filename: str,
         image_page_num: int,
-        user_oid: str,
+        user_oid: Optional[str] = None,
     ) -> Optional[str]:
         raise NotImplementedError("Subclasses must implement this method")
 
     async def download_blob(
-        self, blob_path: str, user_oid: Optional[str] = None, as_bytes: bool = False
-    ) -> Optional[Union[BlobStorageStreamDownloader, AdlsBlobStorageStreamDownloader, bytes]]:
+        self, blob_path: str, user_oid: Optional[str] = None
+    ) -> Optional[tuple[bytes, BlobProperties]]:
         """
         Downloads a blob from Azure Storage.
         If user_oid is provided, it checks if the blob belongs to the user.
@@ -113,7 +116,9 @@ async def download_blob(
             user_oid: The user's object ID (optional)
 
         Returns:
-            bytes: The content of the blob, or None if not found or access denied
+            Optional[tuple[bytes, BlobProperties]]:
+                - A tuple containing the blob content as bytes and the blob properties
+                - None if blob not found or access denied
         """
         raise NotImplementedError("Subclasses must implement this method")
 
@@ -225,7 +230,7 @@ async def upload_document_image(
         image_bytes: bytes,
         image_filename: str,
         image_page_num: int,
-        user_oid: str,
+        user_oid: Optional[str] = None,
     ) -> Optional[str]:
         """
         Uploads an image from a document to ADLS in a directory structure:
@@ -242,6 +247,8 @@ async def upload_document_image(
         Returns:
             str: The URL of the uploaded file, with forward slashes (not URL-encoded)
         """
+        if user_oid is None:
+            raise ValueError("user_oid must be provided for user-specific operations.")
         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)
@@ -252,18 +259,19 @@ async def upload_document_image(
         return unquote(file_client.url)
 
     async def download_blob(
-        self, blob_path: str, user_oid: Optional[str] = None, as_bytes: bool = False
-    ) -> Optional[Union[AdlsBlobStorageStreamDownloader, bytes]]:
+        self, blob_path: str, user_oid: Optional[str] = None
+    ) -> Optional[tuple[bytes, BlobProperties]]:
         """
         Downloads a blob from Azure Data Lake Storage.
 
         Args:
             blob_path: The path to the blob in the format {user_oid}/{document_name}/images/{image_name}
             user_oid: The user's object ID
-            as_bytes: If True, returns the blob as bytes, otherwise returns a stream downloader
 
         Returns:
-            Optional[Union[AdlsBlobStorageStreamDownloader, bytes]]: A stream downloader for the blob, or bytes if as_bytes=True, or None if not found
+            Optional[tuple[bytes, BlobProperties]]:
+                - A tuple containing the blob content as bytes and the blob properties as a dictionary
+                - None if blob not found or access denied
         """
         if user_oid is None:
             logger.warning("user_oid must be provided for Data Lake Storage operations.")
@@ -289,11 +297,17 @@ async def download_blob(
         try:
             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()
-            if as_bytes:
-                return await blob.readall()
-            else:
-                return blob
+            download_response = await file_client.download_file()
+            content = await download_response.readall()
+
+            # Convert FileProperties to our BlobProperties format
+            properties: BlobProperties = {
+                "content_settings": {
+                    "content_type": download_response.properties.get("content_type", "application/octet-stream")
+                }
+            }
+
+            return content, properties
         except ResourceNotFoundError:
             logger.warning(f"Directory or file not found: {directory_path}/{filename}")
             return None
@@ -449,8 +463,23 @@ async def upload_document_image(
         return blob_client.url
 
     async def download_blob(
-        self, blob_path: str, user_oid: Optional[str] = None, as_bytes: bool = False
-    ) -> Optional[Union[BlobStorageStreamDownloader, AdlsBlobStorageStreamDownloader, bytes]]:
+        self, blob_path: str, user_oid: Optional[str] = None
+    ) -> Optional[tuple[bytes, BlobProperties]]:
+        """
+        Downloads a blob from Azure Blob Storage.
+
+        Args:
+            blob_path: The path to the blob in the storage
+            user_oid: Not used in BlobManager, but included for API compatibility
+
+        Returns:
+            Optional[tuple[bytes, BlobProperties]]:
+                - A tuple containing the blob content as bytes and the blob properties
+                - None if blob not found
+
+        Raises:
+            ValueError: If user_oid is provided (not supported for BlobManager)
+        """
         if user_oid is not None:
             raise ValueError(
                 "user_oid is not supported for BlobManager. Use AdlsBlobManager for user-specific operations."
@@ -461,16 +490,33 @@ async def download_blob(
         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:
+            download_response = await blob_client.download_blob()
+            if not download_response.properties:
                 logger.warning(f"No blob exists for {blob_path}")
                 return None
-            if as_bytes:
-                return await blob.readall()
-            else:
-                return blob
+
+            # Get the content as bytes
+            content = await download_response.readall()
+
+            # Convert BlobProperties to our internal BlobProperties format
+            properties: BlobProperties = {
+                "content_settings": {
+                    "content_type": (
+                        download_response.properties.content_settings.content_type
+                        if (
+                            hasattr(download_response.properties, "content_settings")
+                            and download_response.properties.content_settings
+                            and hasattr(download_response.properties.content_settings, "content_type")
+                        )
+                        else "application/octet-stream"
+                    )
+                }
+            }
+
+            return content, properties
         except ResourceNotFoundError:
             logger.warning("Blob not found: %s", blob_path)
             return None

app/backend/prepdocslib/filestrategy.py

@@ -136,7 +136,7 @@ async def run(self):
             await self.search_manager.remove_content()
 
 
-class UploadUserFileStrategy(FileStrategy):
+class UploadUserFileStrategy:
     """
     Strategy for ingesting a file that has already been uploaded to a ADLS2 storage account
     """

app/backend/prepdocslib/integratedvectorizerstrategy.py

@@ -140,6 +140,11 @@ async def create_skillset(self, index_name: str) -> SearchIndexerSkillset:
         return skillset
 
     async def create_multimodal_skillset(self, index_name: str) -> SearchIndexerSkillset:
+        if self.image_embeddings is None:
+            raise ValueError("Image embeddings client must be provided for multimodal skillset creation.")
+        if self.blob_manager.image_container is None:
+            raise ValueError("Blob manager must have an image container set for multimodal skillset creation.")
+
         document_layout_skill = DocumentIntelligenceLayoutSkill(
             description="Layout skill to read documents",
             context="/document",