import json import logging import os from typing import Optional import azure.functions as func from azure.identity import DefaultAzureCredential, get_bearer_token_provider from dotenv import load_dotenv from pydantic import BaseModel, Field from src.components.content_understanding_client import ( AzureContentUnderstandingClient, get_existing_analyzer_ids, ) from src.helpers.common import MeasureRunTime from src.helpers.content_understanding import ( cu_fields_dict_to_markdown, draw_fields_on_imgs, enrich_extracted_cu_fields, ) from src.helpers.data_loading import load_visual_obj_bytes_to_pil_imgs_dict from src.helpers.image import pil_img_to_base64_bytes, resize_img_by_max, rotate_img_pil load_dotenv() token_provider = get_bearer_token_provider( DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default" ) bp_content_understanding_document = func.Blueprint() FUNCTION_ROUTE = "content_understanding_document" # Load environment variables CONTENT_UNDERSTANDING_ENDPOINT = os.getenv("CONTENT_UNDERSTANDING_ENDPOINT") # Load existing analyzer schemas with open("config/content_understanding_schemas.json", "r") as f: CONTENT_UNDERSTANDING_SCHEMAS = json.load(f) cu_client = AzureContentUnderstandingClient( endpoint=CONTENT_UNDERSTANDING_ENDPOINT, azure_ad_token_provider=token_provider, api_version="2024-12-01-preview", enable_face_identification=False, ) # Get list of existing CU analyzers existing_cu_analyzer_ids = get_existing_analyzer_ids(cu_client) class FunctionReponseModel(BaseModel): """ Defines the schema that will be returned by the function. We'll use this to ensure that the response contains the correct values and structure, and to allow a partially filled response to be returned in case of an error. """ success: bool = Field( default=False, description="Indicates whether the pipeline was successful." ) error_text: Optional[str] = Field( default=None, description="If an error occurred, this field will contain the error message.", ) func_time_taken_secs: Optional[float] = Field( default=None, description="The total time taken to process the request." ) enriched_cu_fields: Optional[dict] = Field( default=None, description="The enriched fields result from Content Understanding.", ) formatted_fields_md: Optional[str] = Field( default=None, description="The field outputs in a pre-formatted markdown format.", ) cu_raw_response: Optional[dict] = Field( default=None, description="The raw API response from Content Understanding." ) cu_time_taken_secs: Optional[float] = Field( default=None, description="The time taken to extract the text using Content Understanding.", ) result_imgs_with_bboxes: Optional[dict[int, bytes]] = Field( default=None, description="Dictionary of page images with bounding boxes drawn around the extracted fields.", ) @bp_content_understanding_document.route(route=FUNCTION_ROUTE) def content_understanding_document( req: func.HttpRequest, ) -> func.HttpResponse: """ This function processes a request to extract fields from a PDF document using Azure Content Understanding. If an error occurs at any stage, the function will return a partial response with the error message and the fields that have been populated up to that point. """ logging.info(f"Python HTTP trigger function `{FUNCTION_ROUTE}` received a request.") # Create the object to hold all intermediate and final values. We will progressively update # values as each stage of the pipeline is completed, allowing us to return a partial # response in case of an error at any stage. output_model = FunctionReponseModel(success=False) try: # Create error_text and error_code variables. These will be updated as # we move through the pipeline so that if a step fails, the vars reflect # what has failed. If all steps complete successfully, the vars are # never used. error_text = "An error occurred during processing." error_code = 422 func_timer = MeasureRunTime() func_timer.start() # Check the request body request_json_content = json.loads(req.files["json"].read().decode("utf-8")) analyzer_id = request_json_content.get("analyzer_id", None) file_bytes = req.files["file"].read() file_mime_type = req.files["file"].content_type ### 1. Load the images from the PDF/image input error_text = "An error occurred during image extraction." error_code = 500 doc_page_imgs = load_visual_obj_bytes_to_pil_imgs_dict( file_bytes, file_mime_type, starting_idx=1, pdf_img_dpi=100 ) ### 2. Ensure the analyzer exists error_text = "Invalid Analyzer ID." error_code = 500 # Check if the analyzer already exists when the resource was last checked. global existing_cu_analyzer_ids if analyzer_id not in existing_cu_analyzer_ids: # Refresh the list of existing analyzers existing_cu_analyzer_ids = get_existing_analyzer_ids(cu_client) if analyzer_id not in existing_cu_analyzer_ids: # Analyzer is not available or deployed in the resource. raise KeyError( ( f"Analyzer ID '{analyzer_id}' is not available. " "Ensure that the Analyzer has already been created within the AI services resource." ) ) ### 2. Extract the content using Content Understanding error_text = "An error occurred during Content Understanding extraction." with MeasureRunTime() as cu_timer: response = cu_client.begin_analyze( analyzer_id=analyzer_id, file_bytes=file_bytes, ) cu_result = cu_client.poll_result(response) output_model.cu_raw_response = cu_result output_model.cu_time_taken_secs = cu_timer.time_taken ### 3. Enrich the raw CU API result with additional metadata (normalized polygons, page number etc) error_text = "An error occurred during post-processing." enriched_cu_fields = enrich_extracted_cu_fields( cu_result["result"]["contents"][0] ) output_model.enriched_cu_fields = enriched_cu_fields ### 4. Format the fields in markdown for easy reading output_model.formatted_fields_md = cu_fields_dict_to_markdown( enriched_cu_fields ) ### 5. Draw bounding boxes of the extracted fields on images of the input document # With the locations of the extracted fields now known, we can draw # bounding boxes on the image and to make it easier to digest the output. # Draw the bounding boxes on the image pil_imgs = draw_fields_on_imgs( enriched_fields=enriched_cu_fields, pil_imgs=doc_page_imgs ) # Resize the images to reduce transfer size pil_imgs = { page_num: resize_img_by_max(pil_img, max_height=1000, max_width=1000) for page_num, pil_img in pil_imgs.items() } # Rotate the images to be the correct orientation for page_list_idx, (page_num, pil_img) in enumerate(pil_imgs.items()): pil_imgs[page_num] = rotate_img_pil( pil_img, angle=cu_result["result"]["contents"][0]["pages"][page_list_idx][ "angle" ], ) output_model.result_imgs_with_bboxes = { page_num: pil_img_to_base64_bytes(pil_img) for page_num, pil_img in pil_imgs.items() } ### 5. All steps completed successfully, set success=True and return the final result output_model.success = True output_model.func_time_taken_secs = func_timer.stop() return func.HttpResponse( body=output_model.model_dump_json(), mimetype="application/json", status_code=200, ) except Exception as _e: # If an error occurred at any stage, return the partial response. Update the error_text # field to contain the error message, and ensure success=False. output_model.success = False output_model.error_text = error_text output_model.func_time_taken_secs = func_timer.stop() logging.exception(output_model.error_text) return func.HttpResponse( body=output_model.model_dump_json(), mimetype="application/json", status_code=error_code, )