Store bbox as list of pixel floats, add storage container just for extracted images

Author: pamelafox

app/backend/app.py

@@ -66,8 +66,10 @@
     CONFIG_CHAT_HISTORY_COSMOS_ENABLED,
     CONFIG_CREDENTIAL,
     CONFIG_DEFAULT_REASONING_EFFORT,
+    CONFIG_IMAGE_BLOB_CONTAINER_CLIENT,  # Added this line
     CONFIG_INGESTER,
     CONFIG_LANGUAGE_PICKER_ENABLED,
+    CONFIG_MULTIMODAL_ENABLED,
     CONFIG_OPENAI_CLIENT,
     CONFIG_QUERY_REWRITING_ENABLED,
     CONFIG_REASONING_EFFORT_ENABLED,
@@ -84,7 +86,6 @@
     CONFIG_USER_BLOB_CONTAINER_CLIENT,
     CONFIG_USER_UPLOAD_ENABLED,
     CONFIG_VECTOR_SEARCH_ENABLED,
-    CONFIG_MULTIMODAL_ENABLED
 )
 from core.authentication import AuthenticationHelper
 from core.sessionhelper import create_session_id
@@ -182,7 +183,9 @@ async def ask(auth_claims: dict[str, Any]):
     context["auth_claims"] = auth_claims
     try:
         approach: Approach = cast(Approach, current_app.config[CONFIG_ASK_APPROACH])
-        r = await approach.run(request_json["messages"], context=context, session_state=request_json.get("session_state"))
+        r = await approach.run(
+            request_json["messages"], context=context, session_state=request_json.get("session_state")
+        )
         return jsonify(r)
     except Exception as error:
         return error_response(error, "/ask")
@@ -404,6 +407,7 @@ async def setup_clients():
     # Replace these with your own values, either in environment variables or directly here
     AZURE_STORAGE_ACCOUNT = os.environ["AZURE_STORAGE_ACCOUNT"]
     AZURE_STORAGE_CONTAINER = os.environ["AZURE_STORAGE_CONTAINER"]
+    AZURE_IMAGESTORAGE_CONTAINER = os.environ.get("AZURE_IMAGESTORAGE_CONTAINER")
     AZURE_USERSTORAGE_ACCOUNT = os.environ.get("AZURE_USERSTORAGE_ACCOUNT")
     AZURE_USERSTORAGE_CONTAINER = os.environ.get("AZURE_USERSTORAGE_CONTAINER")
     AZURE_SEARCH_SERVICE = os.environ["AZURE_SEARCH_SERVICE"]
@@ -511,6 +515,15 @@ async def setup_clients():
         f"https://{AZURE_STORAGE_ACCOUNT}.blob.core.windows.net", AZURE_STORAGE_CONTAINER, credential=azure_credential
     )
 
+    # Set up the image storage container client if configured
+    image_blob_container_client = None
+    if AZURE_IMAGESTORAGE_CONTAINER:
+        image_blob_container_client = ContainerClient(
+            f"https://{AZURE_STORAGE_ACCOUNT}.blob.core.windows.net",
+            AZURE_IMAGESTORAGE_CONTAINER,
+            credential=azure_credential,
+        )
+
     # Set up authentication helper
     search_index = None
     if AZURE_USE_AUTHENTICATION:
@@ -636,6 +649,8 @@ async def setup_clients():
     current_app.config[CONFIG_SEARCH_CLIENT] = search_client
     current_app.config[CONFIG_AGENT_CLIENT] = agent_client
     current_app.config[CONFIG_BLOB_CONTAINER_CLIENT] = blob_container_client
+    if image_blob_container_client:
+        current_app.config[CONFIG_IMAGE_BLOB_CONTAINER_CLIENT] = image_blob_container_client
     current_app.config[CONFIG_AUTH_CLIENT] = auth_helper
 
     current_app.config[CONFIG_SEMANTIC_RANKER_DEPLOYED] = AZURE_SEARCH_SEMANTIC_RANKER != "disabled"
@@ -710,12 +725,15 @@ async def setup_clients():
         vision_token_provider=token_provider,
     )
 
+
 @bp.after_app_serving
 async def close_clients():
     await current_app.config[CONFIG_SEARCH_CLIENT].close()
     await current_app.config[CONFIG_BLOB_CONTAINER_CLIENT].close()
     if current_app.config.get(CONFIG_USER_BLOB_CONTAINER_CLIENT):
         await current_app.config[CONFIG_USER_BLOB_CONTAINER_CLIENT].close()
+    if current_app.config.get(CONFIG_IMAGE_BLOB_CONTAINER_CLIENT):
+        await current_app.config[CONFIG_IMAGE_BLOB_CONTAINER_CLIENT].close()
 
 
 def create_app():

app/backend/config.py

@@ -3,6 +3,7 @@
 CONFIG_ASK_APPROACH = "ask_approach"
 CONFIG_CHAT_APPROACH = "chat_approach"
 CONFIG_BLOB_CONTAINER_CLIENT = "blob_container_client"
+CONFIG_IMAGE_BLOB_CONTAINER_CLIENT = "image_blob_container_client"
 CONFIG_USER_UPLOAD_ENABLED = "user_upload_enabled"
 CONFIG_USER_BLOB_CONTAINER_CLIENT = "user_blob_container_client"
 CONFIG_AUTH_CLIENT = "auth_client"
@@ -30,4 +31,4 @@
 CONFIG_COSMOS_HISTORY_CLIENT = "cosmos_history_client"
 CONFIG_COSMOS_HISTORY_CONTAINER = "cosmos_history_container"
 CONFIG_COSMOS_HISTORY_VERSION = "cosmos_history_version"
-CONFIG_MULTIMODAL_ENABLED = "multimodal_enabled"
+CONFIG_MULTIMODAL_ENABLED = "multimodal_enabled"

app/backend/prepdocs.py

@@ -2,13 +2,14 @@
 import asyncio
 import logging
 import os
+from enum import Enum
 from typing import Optional, Union
 
 from azure.core.credentials import AzureKeyCredential
 from azure.core.credentials_async import AsyncTokenCredential
 from azure.identity.aio import AzureDeveloperCliCredential, get_bearer_token_provider
-from rich.logging import RichHandler
 from openai import AsyncAzureOpenAI, AsyncOpenAI
+from rich.logging import RichHandler
 
 from load_azd_env import load_azd_env
 from prepdocslib.blobmanager import BlobManager
@@ -31,11 +32,14 @@
     LocalListFileStrategy,
 )
 from prepdocslib.parser import Parser
-from prepdocslib.pdfparser import DocumentAnalysisParser, LocalPdfParser, MediaDescriptionStrategy
+from prepdocslib.pdfparser import (
+    DocumentAnalysisParser,
+    LocalPdfParser,
+    MediaDescriptionStrategy,
+)
 from prepdocslib.strategy import DocumentAction, SearchInfo, Strategy
 from prepdocslib.textparser import TextParser
 from prepdocslib.textsplitter import SentenceTextSplitter, SimpleTextSplitter
-from enum import Enum
 
 logger = logging.getLogger("scripts")
 
@@ -86,11 +90,14 @@ def setup_blob_manager(
     subscription_id: str,
     store_page_images: bool,
     storage_key: Union[str, None] = None,
+    image_storage_container: Union[str, None] = None,  # Added this parameter
 ):
     storage_creds: Union[AsyncTokenCredential, str] = azure_credential if storage_key is None else storage_key
+
     return BlobManager(
         endpoint=f"https://{storage_account}.blob.core.windows.net",
         container=storage_container,
+        image_container=image_storage_container,
         account=storage_account,
         credential=storage_creds,
         resourceGroup=storage_resource_group,
@@ -178,6 +185,7 @@ def setup_embeddings_service(
             disable_batch=disable_batch_vectors,
         )
 
+
 def setup_openai_client(
     openai_host: OpenAIHost,
     azure_openai_api_key: Union[str, None] = None,
@@ -231,6 +239,7 @@ def setup_openai_client(
         )
     return openai_client
 
+
 def setup_file_processors(
     azure_credential: AsyncTokenCredential,
     document_intelligence_service: Union[str, None],
@@ -255,7 +264,15 @@ def setup_file_processors(
         doc_int_parser = DocumentAnalysisParser(
             endpoint=f"https://{document_intelligence_service}.cognitiveservices.azure.com/",
             credential=documentintelligence_creds,
-            media_description_strategy = MediaDescriptionStrategy.OPENAI if use_multimodal else MediaDescriptionStrategy.CONTENTUNDERSTANDING if use_content_understanding else MediaDescriptionStrategy.NONE,
+            media_description_strategy=(
+                MediaDescriptionStrategy.OPENAI
+                if use_multimodal
+                else (
+                    MediaDescriptionStrategy.CONTENTUNDERSTANDING
+                    if use_content_understanding
+                    else MediaDescriptionStrategy.NONE
+                )
+            ),
             openai_client=openai_client,
             openai_model=openai_model,
             openai_deployment=openai_deployment,
@@ -384,7 +401,9 @@ async def main(strategy: Strategy, setup_index: bool = True):
     args = parser.parse_args()
 
     if args.verbose:
-        logging.basicConfig(format="%(message)s", datefmt="[%X]", handlers=[RichHandler(rich_tracebacks=True)], level=logging.WARNING)
+        logging.basicConfig(
+            format="%(message)s", datefmt="[%X]", handlers=[RichHandler(rich_tracebacks=True)], level=logging.WARNING
+        )
         # We only set the level to INFO for our logger,
         # to avoid seeing the noisy INFO level logs from the Azure SDKs
         logger.setLevel(logging.DEBUG)
@@ -448,6 +467,7 @@ async def main(strategy: Strategy, setup_index: bool = True):
         subscription_id=os.environ["AZURE_SUBSCRIPTION_ID"],
         store_page_images=use_multimodal,
         storage_key=clean_key_if_exists(args.storagekey),
+        image_storage_container=os.environ.get("AZURE_IMAGESTORAGE_CONTAINER"),  # Pass the image container
     )
     list_file_strategy = setup_list_file_strategy(
         azure_credential=azd_credential,
@@ -460,7 +480,7 @@ async def main(strategy: Strategy, setup_index: bool = True):
 
     openai_host = OpenAIHost(os.environ["OPENAI_HOST"])
     # 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-06-01"
+    azure_openai_api_version = os.getenv("AZURE_OPENAI_API_VERSION") or "2024-06-01"
     emb_model_dimensions = 1536
     if os.getenv("AZURE_OPENAI_EMB_DIMENSIONS"):
         emb_model_dimensions = int(os.environ["AZURE_OPENAI_EMB_DIMENSIONS"])
@@ -490,7 +510,6 @@ async def main(strategy: Strategy, setup_index: bool = True):
         openai_organization=os.getenv("OPENAI_ORGANIZATION"),
     )
 
-
     ingestion_strategy: Strategy
     if use_int_vectorization:
 

app/backend/prepdocslib/blobmanager.py

@@ -35,11 +35,13 @@ def __init__(
         resourceGroup: str,
         subscriptionId: str,
         store_page_images: bool = False,
+        image_container: Optional[str] = None,  # Added this parameter
     ):
         self.endpoint = endpoint
         self.credential = credential
         self.account = account
         self.container = container
+        self.image_container = image_container
         self.store_page_images = store_page_images
         self.resourceGroup = resourceGroup
         self.subscriptionId = subscriptionId
@@ -60,11 +62,17 @@ 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 None
-    
-    async def upload_document_image(self, document_file: File, image_bytes: bytes, image_filename: str) -> Optional[str]:
+
+    async def upload_document_image(
+        self, document_file: File, image_bytes: bytes, image_filename: str
+    ) -> Optional[str]:
+        if self.image_container is None:
+            raise ValueError(
+                "Image container name is not set. Re-run `azd provision` to automatically set up the images container."
+            )
         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:
+        ) 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()
             blob_name = BlobManager.blob_name_from_file_name(document_file.content.name) + "/" + image_filename

app/backend/prepdocslib/page.py

@@ -1,16 +1,16 @@
-from typing import Sequence
 from dataclasses import dataclass, field
 
 
 @dataclass
 class ImageOnPage:
     bytes: bytes
-    bbox: tuple[float, float, float, float]
+    bbox: list[float, float, float, float]  # Pixels
     filename: str
     description: str
     url: str | None = None
     embedding: list[float] | None = None
 
+
 @dataclass
 class Page:
     """
@@ -21,11 +21,13 @@ class Page:
         offset (int): If the text of the entire Document was concatenated into a single string, the index of the first character on the page. For example, if page 1 had the text "hello" and page 2 had the text "world", the offset of page 2 is 5 ("hellow")
         text (str): The text of the page
     """
+
     page_num: int
     offset: int
     text: str
     images: list[ImageOnPage] = field(default_factory=list)
 
+
 @dataclass
 class SplitPage:
     """
@@ -35,6 +37,7 @@ class SplitPage:
         page_num (int): Page number (0-indexed)
         text (str): The text of the section
     """
+
     page_num: int
     text: str
-    images: list[ImageOnPage] = field(default_factory=list)
+    images: list[ImageOnPage] = field(default_factory=list)