Fix media description with OpenAI

pamelafox · pamelafox · commit 7c8f8255fa4f · 2025-06-02T20:32:37.000Z
diff --git a/app/backend/app.py b/app/backend/app.py
@@ -84,6 +84,7 @@
     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
@@ -705,6 +706,8 @@ async def setup_clients():
         query_speller=AZURE_SEARCH_QUERY_SPELLER,
         prompt_manager=prompt_manager,
         reasoning_effort=OPENAI_REASONING_EFFORT,
+        vision_endpoint=AZURE_VISION_ENDPOINT,
+        vision_token_provider=token_provider,
     )
 
 @bp.after_app_serving
@@ -734,7 +737,7 @@ def create_app():
 
     # Log levels should be one of https://docs.python.org/3/library/logging.html#logging-levels
     # Set root level to WARNING to avoid seeing overly verbose logs from SDKS
-    logging.basicConfig(level=logging.WARNING)
+    logging.basicConfig(level=logging.DEBUG)
     # Set our own logger levels to INFO by default
     app_level = os.getenv("APP_LOG_LEVEL", "INFO")
     app.logger.setLevel(os.getenv("APP_LOG_LEVEL", app_level))
diff --git a/app/backend/approaches/chatreadretrieveread.py b/app/backend/approaches/chatreadretrieveread.py
@@ -1,5 +1,5 @@
 from collections.abc import Awaitable
-from typing import Any, Optional, Union, cast
+from typing import Any, Optional, Union, cast, Callable
 
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
 from azure.search.documents.aio import SearchClient
@@ -47,6 +47,8 @@ def __init__(
         query_speller: str,
         prompt_manager: PromptManager,
         reasoning_effort: Optional[str] = None,
+        vision_endpoint: Optional[str] = None,
+        vision_token_provider: Callable[[], Awaitable[str]],
     ):
         self.search_client = search_client
         self.search_index_name = search_index_name
@@ -71,6 +73,8 @@ def __init__(
         self.answer_prompt = self.prompt_manager.load_prompt("chat_answer_question.prompty")
         self.reasoning_effort = reasoning_effort
         self.include_token_usage = True
+        self.vision_endpoint = vision_endpoint
+        self.vision_token_provider = vision_token_provider
 
     async def run_until_final_call(
         self,
diff --git a/app/backend/prepdocs.py b/app/backend/prepdocs.py
@@ -35,6 +35,7 @@
 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")
 
@@ -126,15 +127,23 @@ def setup_list_file_strategy(
     return list_file_strategy
 
 
+class OpenAIHost(str, Enum):
+    OPENAI = "openai"
+    AZURE = "azure"
+    AZURE_CUSTOM = "azure_custom"
+    LOCAL = "local"
+
+
 def setup_embeddings_service(
     azure_credential: AsyncTokenCredential,
-    openai_host: str,
-    openai_model_name: str,
-    openai_service: Union[str, None],
-    openai_custom_url: Union[str, None],
-    openai_deployment: Union[str, None],
-    openai_dimensions: int,
-    openai_api_version: str,
+    openai_host: OpenAIHost,
+    emb_model_name: str,
+    emb_model_dimensions: int,
+    azure_openai_service: Union[str, None],
+    azure_openai_custom_url: Union[str, None],
+    azure_openai_deployment: Union[str, None],
+    azure_openai_key: Union[str, None],
+    azure_openai_api_version: str,
     openai_key: Union[str, None],
     openai_org: Union[str, None],
     disable_vectors: bool = False,
@@ -144,31 +153,83 @@ def setup_embeddings_service(
         logger.info("Not setting up embeddings service")
         return None
 
-    if openai_host != "openai":
+    if openai_host in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]:
         azure_open_ai_credential: Union[AsyncTokenCredential, AzureKeyCredential] = (
-            azure_credential if openai_key is None else AzureKeyCredential(openai_key)
+            azure_credential if azure_openai_key is None else AzureKeyCredential(azure_openai_key)
         )
         return AzureOpenAIEmbeddingService(
-            open_ai_service=openai_service,
-            open_ai_custom_url=openai_custom_url,
-            open_ai_deployment=openai_deployment,
-            open_ai_model_name=openai_model_name,
-            open_ai_dimensions=openai_dimensions,
-            open_ai_api_version=openai_api_version,
+            open_ai_service=azure_openai_service,
+            open_ai_custom_url=azure_openai_custom_url,
+            open_ai_deployment=azure_openai_deployment,
+            open_ai_model_name=emb_model_name,
+            open_ai_dimensions=emb_model_dimensions,
+            open_ai_api_version=azure_openai_api_version,
             credential=azure_open_ai_credential,
             disable_batch=disable_batch_vectors,
         )
     else:
         if openai_key is None:
             raise ValueError("OpenAI key is required when using the non-Azure OpenAI API")
         return OpenAIEmbeddingService(
-            open_ai_model_name=openai_model_name,
-            open_ai_dimensions=openai_dimensions,
+            open_ai_model_name=emb_model_name,
+            open_ai_dimensions=emb_model_dimensions,
             credential=openai_key,
             organization=openai_org,
             disable_batch=disable_batch_vectors,
         )
 
+def setup_openai_client(
+    openai_host: OpenAIHost,
+    azure_openai_api_key: Union[str, None] = None,
+    azure_openai_api_version: Union[str, None] = None,
+    azure_openai_service: Union[str, None] = None,
+    azure_openai_custom_url: Union[str, None] = None,
+    azure_credential: AsyncTokenCredential = None,
+    openai_api_key: Union[str, None] = None,
+    openai_organization: Union[str, None] = None,
+):
+    if openai_host not in OpenAIHost:
+        raise ValueError(f"Invalid OPENAI_HOST value: {openai_host}. Must be one of {[h.value for h in OpenAIHost]}.")
+
+    if openai_host in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]:
+        if openai_host == OpenAIHost.AZURE_CUSTOM:
+            logger.info("OPENAI_HOST is azure_custom, setting up Azure OpenAI custom client")
+            if not azure_openai_custom_url:
+                raise ValueError("AZURE_OPENAI_CUSTOM_URL must be set when OPENAI_HOST is azure_custom")
+            endpoint = azure_openai_custom_url
+        else:
+            logger.info("OPENAI_HOST is azure, setting up Azure OpenAI client")
+            if not azure_openai_service:
+                raise ValueError("AZURE_OPENAI_SERVICE must be set when OPENAI_HOST is azure")
+            endpoint = f"https://{azure_openai_service}.openai.azure.com"
+        if azure_openai_api_key:
+            logger.info("AZURE_OPENAI_API_KEY_OVERRIDE found, using as api_key for Azure OpenAI client")
+            openai_client = AsyncAzureOpenAI(
+                api_version=azure_openai_api_version, azure_endpoint=endpoint, api_key=azure_openai_api_key
+            )
+        else:
+            logger.info("Using Azure credential (passwordless authentication) for Azure OpenAI client")
+            token_provider = get_bearer_token_provider(azure_credential, "https://cognitiveservices.azure.com/.default")
+            openai_client = AsyncAzureOpenAI(
+                api_version=azure_openai_api_version,
+                azure_endpoint=endpoint,
+                azure_ad_token_provider=token_provider,
+            )
+    elif openai_host == OpenAIHost.LOCAL:
+        logger.info("OPENAI_HOST is local, setting up local OpenAI client for OPENAI_BASE_URL with no key")
+        openai_client = AsyncOpenAI(
+            base_url=os.environ["OPENAI_BASE_URL"],
+            api_key="no-key-required",
+        )
+    else:
+        logger.info(
+            "OPENAI_HOST is not azure, setting up OpenAI client using OPENAI_API_KEY and OPENAI_ORGANIZATION environment variables"
+        )
+        openai_client = AsyncOpenAI(
+            api_key=openai_api_key,
+            organization=openai_organization,
+        )
+    return openai_client
 
 def setup_file_processors(
     azure_credential: AsyncTokenCredential,
@@ -194,7 +255,7 @@ def setup_file_processors(
         doc_int_parser = DocumentAnalysisParser(
             endpoint=f"https://{document_intelligence_service}.cognitiveservices.azure.com/",
             credential=documentintelligence_creds,
-            media_description_strategy = "openai" if use_multimodal else "contentunderstanding" if use_content_understanding else "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,
@@ -323,7 +384,7 @@ 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)])
+        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)
@@ -397,31 +458,38 @@ async def main(strategy: Strategy, setup_index: bool = True):
         datalake_key=clean_key_if_exists(args.datalakekey),
     )
 
-    openai_host = os.environ["OPENAI_HOST"]
-    openai_key = None
-    if os.getenv("AZURE_OPENAI_API_KEY_OVERRIDE"):
-        openai_key = os.getenv("AZURE_OPENAI_API_KEY_OVERRIDE")
-    elif not openai_host.startswith("azure") and os.getenv("OPENAI_API_KEY"):
-        openai_key = os.getenv("OPENAI_API_KEY")
-
-    openai_dimensions = 1536
+    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"
+    emb_model_dimensions = 1536
     if os.getenv("AZURE_OPENAI_EMB_DIMENSIONS"):
-        openai_dimensions = int(os.environ["AZURE_OPENAI_EMB_DIMENSIONS"])
+        emb_model_dimensions = int(os.environ["AZURE_OPENAI_EMB_DIMENSIONS"])
     openai_embeddings_service = setup_embeddings_service(
         azure_credential=azd_credential,
         openai_host=openai_host,
-        openai_model_name=os.environ["AZURE_OPENAI_EMB_MODEL_NAME"],
-        openai_service=os.getenv("AZURE_OPENAI_SERVICE"),
-        openai_custom_url=os.getenv("AZURE_OPENAI_CUSTOM_URL"),
-        openai_deployment=os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT"),
-        # https://learn.microsoft.com/azure/ai-services/openai/api-version-deprecation#latest-ga-api-release
-        openai_api_version=os.getenv("AZURE_OPENAI_API_VERSION") or "2024-06-01",
-        openai_dimensions=openai_dimensions,
-        openai_key=clean_key_if_exists(openai_key),
+        emb_model_name=os.environ["AZURE_OPENAI_EMB_MODEL_NAME"],
+        emb_model_dimensions=emb_model_dimensions,
+        azure_openai_service=os.getenv("AZURE_OPENAI_SERVICE"),
+        azure_openai_custom_url=os.getenv("AZURE_OPENAI_CUSTOM_URL"),
+        azure_openai_deployment=os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT"),
+        azure_openai_api_version=azure_openai_api_version,
+        azure_openai_key=os.getenv("AZURE_OPENAI_API_KEY_OVERRIDE"),
+        openai_key=clean_key_if_exists(os.getenv("OPENAI_API_KEY")),
         openai_org=os.getenv("OPENAI_ORGANIZATION"),
         disable_vectors=dont_use_vectors,
         disable_batch_vectors=args.disablebatchvectors,
     )
+    openai_client = setup_openai_client(
+        openai_host=openai_host,
+        azure_openai_api_version=azure_openai_api_version,
+        azure_openai_service=os.getenv("AZURE_OPENAI_SERVICE"),
+        azure_openai_custom_url=os.getenv("AZURE_OPENAI_CUSTOM_URL"),
+        azure_openai_api_key=os.getenv("AZURE_OPENAI_API_KEY_OVERRIDE"),
+        azure_credential=azd_credential,
+        openai_api_key=clean_key_if_exists(os.getenv("OPENAI_API_KEY")),
+        openai_organization=os.getenv("OPENAI_ORGANIZATION"),
+    )
+
 
     ingestion_strategy: Strategy
     if use_int_vectorization:
@@ -452,6 +520,9 @@ async def main(strategy: Strategy, setup_index: bool = True):
             use_content_understanding=use_content_understanding,
             use_multimodal=use_multimodal,
             content_understanding_endpoint=os.getenv("AZURE_CONTENTUNDERSTANDING_ENDPOINT"),
+            openai_client=openai_client,
+            openai_model=os.getenv("AZURE_OPENAI_CHATGPT_MODEL"),
+            openai_deployment=os.getenv("AZURE_OPENAI_CHATGPT_DEPLOYMENT") if openai_host == OpenAIHost.AZURE else None,
         )
 
         image_embeddings_service = setup_image_embeddings_service(
diff --git a/app/backend/prepdocslib/blobmanager.py b/app/backend/prepdocslib/blobmanager.py
@@ -60,11 +60,11 @@ 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
 
-            if self.store_page_images:
-                if os.path.splitext(file.content.name)[1].lower() == ".pdf":
-                    return await self.upload_pdf_blob_images(service_client, container_client, file)
-                else:
-                    logger.info("File %s is not a PDF, skipping image upload", file.content.name)
+            #if self.store_page_images:
+            #    if os.path.splitext(file.content.name)[1].lower() == ".pdf":
+            #        return await self.upload_pdf_blob_images(service_client, container_client, file)
+            #    else:
+            #        logger.info("File %s is not a PDF, skipping image upload", file.content.name)
 
         return None
 
diff --git a/app/backend/prepdocslib/mediadescriber.py b/app/backend/prepdocslib/mediadescriber.py
@@ -121,17 +121,20 @@ async def describe_image(self, image_bytes: bytes) -> str:
 
         response = await self.openai_client.chat.completions.create(
             model=self.model if self.deployment is None else self.deployment,
+            max_tokens=500,
             messages=[
                 {
                     "role": "system",
-                    "content": "You are a helpful assistant that describes images.",
+                    "content": "You are a helpful assistant that describes images from organizational documents.",
                 },
                 {
                     "role": "user",
                     "content": 
-                    [{"text": "Describe this image in detail", "type": "text"},
-                    {"image_url": {"url": image_datauri}, "type": "image_url"}]
+                    [{"text": "Describe image with no more than 5 sentences. Do not speculate about anything you don't know.", "type": "text"},
+                    {"image_url": {"url": image_datauri}, "type": "image_url", "detail": "low"}]
                 }
             ])
-        return response.choices[0].message.content.strip() if response.choices else ""
+        description = response.choices[0].message.content.strip() if response.choices else ""
+        print(description)
+        return description
 
diff --git a/app/backend/prepdocslib/pdfparser.py b/app/backend/prepdocslib/pdfparser.py
@@ -61,7 +61,6 @@ def __init__(
         endpoint: str,
         credential: Union[AsyncTokenCredential, AzureKeyCredential],
         model_id="prebuilt-layout",
-        include_media_description: bool = False,
         media_description_strategy: Enum = MediaDescriptionStrategy.NONE,
         # If using OpenAI, this is the client to use
         openai_client: Union[AsyncOpenAI, None] = None,
@@ -275,6 +274,10 @@ def crop_image_from_pdf_page(
         pix = page.get_pixmap(matrix=pymupdf.Matrix(page_dpi / bbox_dpi, page_dpi / bbox_dpi), clip=rect)
 
         img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
+        # print out the number of pixels
+        print(f"Cropped image size: {img.size} pixels")
         bytes_io = io.BytesIO()
         img.save(bytes_io, format="PNG")
+        with open(f"cropped_page_{page_number + 1}.png", "wb") as f:
+            f.write(bytes_io.getvalue())
         return bytes_io.getvalue()
diff --git a/app/backend/prepdocslib/searchmanager.py b/app/backend/prepdocslib/searchmanager.py
@@ -298,6 +298,7 @@ async def create_index(self):
                     field.name == self.field_name_embedding for field in existing_index.fields
                 ):
                     logger.info("Adding %s field for text embeddings", self.field_name_embedding)
+                    embedding_field.stored = True
                     existing_index.fields.append(embedding_field)
                     if existing_index.vector_search is None:
                         raise ValueError("Vector search is not enabled for the existing index")
diff --git a/app/frontend/package-lock.json b/app/frontend/package-lock.json
