Mypy fixes

Author: pamelafox

app/backend/app.py

@@ -397,7 +397,7 @@ async def setup_clients():
     AZURE_SEARCH_INDEX = os.environ["AZURE_SEARCH_INDEX"]
     AZURE_SEARCH_AGENT = os.getenv("AZURE_SEARCH_AGENT", "")
     # Shared by all OpenAI deployments
-    OPENAI_HOST = os.getenv("OPENAI_HOST", "azure")
+    OPENAI_HOST = OpenAIHost(os.getenv("OPENAI_HOST", "azure"))
     OPENAI_CHATGPT_MODEL = os.environ["AZURE_OPENAI_CHATGPT_MODEL"]
     AZURE_OPENAI_SEARCHAGENT_MODEL = os.getenv("AZURE_OPENAI_SEARCHAGENT_MODEL")
     AZURE_OPENAI_SEARCHAGENT_DEPLOYMENT = os.getenv("AZURE_OPENAI_SEARCHAGENT_DEPLOYMENT")
@@ -407,9 +407,13 @@ async def setup_clients():
     # Used with Azure OpenAI deployments
     AZURE_OPENAI_SERVICE = os.getenv("AZURE_OPENAI_SERVICE")
     AZURE_OPENAI_CHATGPT_DEPLOYMENT = (
-        os.getenv("AZURE_OPENAI_CHATGPT_DEPLOYMENT") if OPENAI_HOST.startswith("azure") else None
+        os.getenv("AZURE_OPENAI_CHATGPT_DEPLOYMENT")
+        if OPENAI_HOST in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]
+        else None
+    )
+    AZURE_OPENAI_EMB_DEPLOYMENT = (
+        os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT") if OPENAI_HOST in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM] else None
     )
-    AZURE_OPENAI_EMB_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT") if OPENAI_HOST.startswith("azure") else None
     AZURE_OPENAI_CUSTOM_URL = os.getenv("AZURE_OPENAI_CUSTOM_URL")
     # https://learn.microsoft.com/azure/ai-services/openai/api-version-deprecation#latest-ga-api-release
     AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION") or "2024-10-21"

app/backend/approaches/approach.py

@@ -449,6 +449,8 @@ class ExtraArgs(TypedDict, total=False):
         return VectorizedQuery(vector=query_vector, k_nearest_neighbors=50, fields=self.embedding_field)
 
     async def compute_multimodal_embedding(self, q: str):
+        if not self.image_embeddings_client:
+            raise ValueError("Approach is missing an image embeddings client for multimodal queries")
         multimodal_query_vector = await self.image_embeddings_client.create_embedding_for_text(q)
         return VectorizedQuery(vector=multimodal_query_vector, k_nearest_neighbors=50, fields="images/embedding")
 

app/backend/prepdocslib/blobmanager.py

@@ -12,11 +12,10 @@
 from azure.storage.blob.aio import (
     StorageStreamDownloader as BlobStorageStreamDownloader,
 )
-from azure.storage.filedatalake import DataLakeDirectoryClient
-from azure.storage.filedatalake import (
+from azure.storage.filedatalake.aio import DataLakeDirectoryClient, FileSystemClient
+from azure.storage.filedatalake.aio import (
     StorageStreamDownloader as AdlsBlobStorageStreamDownloader,
 )
-from azure.storage.filedatalake.aio import FileSystemClient
 from PIL import Image, ImageDraw, ImageFont
 
 from .listfilestrategy import File
@@ -254,16 +253,17 @@ async def upload_document_image(
 
     async def download_blob(
         self, blob_path: str, user_oid: Optional[str] = None, as_bytes: bool = False
-    ) -> Optional[Union[AdlsBlobStorageStreamDownloader, BlobStorageStreamDownloader]]:
+    ) -> Optional[Union[AdlsBlobStorageStreamDownloader, bytes]]:
         """
         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, BlobStorageStreamDownloader]]: A stream downloader for the blob, or None if not found
+            Optional[Union[AdlsBlobStorageStreamDownloader, bytes]]: A stream downloader for the blob, or bytes if as_bytes=True, or None if not found
         """
         if user_oid is None:
             logger.warning("user_oid must be provided for Data Lake Storage operations.")
@@ -322,8 +322,8 @@ async def remove_blob(self, filename: str, user_oid: str) -> None:
         await file_client.delete_file()
 
         # Try to delete any associated image directories
+        image_directory_path = self._get_image_directory_path(filename, user_oid)
         try:
-            image_directory_path = self._get_image_directory_path(filename, user_oid)
             image_directory_client = await self._ensure_directory(
                 directory_path=image_directory_path, user_oid=user_oid
             )
@@ -408,7 +408,7 @@ def get_managedidentity_connectionstring(self):
             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 upload_blob(self, file: File) -> Optional[list[str]]:
+    async def upload_blob(self, file: File) -> str:
         container_client = self.blob_service_client.get_container_client(self.container)
         if not await container_client.exists():
             await container_client.create_container()
@@ -421,6 +421,8 @@ async def upload_blob(self, file: File) -> Optional[list[str]]:
                 blob_client = await container_client.upload_blob(blob_name, reopened_file, overwrite=True)
                 file.url = blob_client.url
 
+        return unquote(file.url)
+
     async def upload_document_image(
         self,
         document_filename: str,

app/backend/prepdocslib/filestrategy.py

@@ -3,7 +3,7 @@
 
 from azure.core.credentials import AzureKeyCredential
 
-from .blobmanager import BlobManager
+from .blobmanager import AdlsBlobManager, BaseBlobManager, BlobManager
 from .embeddings import ImageEmbeddings, OpenAIEmbeddings
 from .fileprocessor import FileProcessor
 from .listfilestrategy import File, ListFileStrategy
@@ -18,7 +18,7 @@ async def parse_file(
     file: File,
     file_processors: dict[str, FileProcessor],
     category: Optional[str] = None,
-    blob_manager: Optional[BlobManager] = None,
+    blob_manager: Optional[BaseBlobManager] = None,
     image_embeddings_client: Optional[ImageEmbeddings] = None,
     user_oid: Optional[str] = None,
 ) -> list[Section]:
@@ -145,10 +145,10 @@ def __init__(
         self,
         search_info: SearchInfo,
         file_processors: dict[str, FileProcessor],
+        blob_manager: AdlsBlobManager,
+        search_field_name_embedding: Optional[str] = None,
         embeddings: Optional[OpenAIEmbeddings] = None,
         image_embeddings: Optional[ImageEmbeddings] = None,
-        search_field_name_embedding: Optional[str] = None,
-        blob_manager: Optional[BlobManager] = None,
     ):
         self.file_processors = file_processors
         self.embeddings = embeddings