Extract QA Audit Findings from PDF Reports with Azure Document Intelligence SDK and Haystack

"""
Production implementation for extracting QA audit findings from PDF reports using Azure Document Intelligence and Haystack.
This module integrates Azure's intelligence capabilities with Haystack for efficient data retrieval.
"""

from typing import Dict, Any, List
import os
import logging
import time
from azure.ai.formrecognizer import DocumentAnalysisClient
from azure.core.credentials import AzureKeyCredential

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Config:
    """
    Configuration class to hold environment variables.
    """
    endpoint: str = os.getenv('AZURE_FORM_RECOGNIZER_ENDPOINT')
    api_key: str = os.getenv('AZURE_FORM_RECOGNIZER_API_KEY')

def validate_input(file_path: str) -> None:
    """Validate the input PDF file path.
    
    Args:
        file_path: The file path to validate
    Raises:
        ValueError: If the file path is invalid or the file does not exist
    """
    if not os.path.isfile(file_path):
        raise ValueError(f'Invalid file path: {file_path}')
    logger.info(f'Input file validated: {file_path}')  # Log successful validation

def sanitize_fields(data: Dict[str, Any]) -> Dict[str, Any]:
    """Sanitize the extracted data fields to prevent injection.
    
    Args:
        data: Extracted data from Azure
    Returns:
        Sanitized data dictionary
    """
    sanitized_data = {k: str(v).strip() for k, v in data.items()}
    logger.debug(f'Sanitized data: {sanitized_data}')  # Log sanitized data
    return sanitized_data

def fetch_data(file_path: str) -> Dict[str, Any]:
    """Fetch PDF data using Azure Document Intelligence.
    
    Args:
        file_path: Path to the PDF file
    Returns:
        The extracted data from the PDF
    Raises:
        Exception: If an error occurs during document analysis
    """
    try:
        client = DocumentAnalysisClient(endpoint=Config.endpoint,
                                         credential=AzureKeyCredential(Config.api_key))
        with open(file_path, "rb") as fd:
            poller = client.begin_read_in_stream(fd, 1, language='en')
        result = poller.result()
        logger.info('Document processed successfully.')  # Log success
        return result
    except Exception as e:
        logger.error(f'Error fetching data: {e}')  # Log error
        raise

def transform_records(raw_data: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Transform the raw data into a structured format.
    
    Args:
        raw_data: The data extracted from Azure
    Returns:
        A list of structured records
    """
    records = []
    for page_result in raw_data.pages:
        for line in page_result.lines:
            record = {
                'text': line.text,
                'bounding_box': line.bounding_box
            }
            records.append(sanitize_fields(record))
    logger.info(f'Transformed records: {len(records)} found.')  # Log number of records
    return records

def save_to_db(records: List[Dict[str, Any]]) -> None:
    """Save processed records to a database (mock implementation).
    
    Args:
        records: List of records to save
    """
    # This is a placeholder for actual database save logic
    logger.info(f'Saving {len(records)} records to the database.')  # Log save action

def handle_errors(e: Exception) -> None:
    """Handle exceptions gracefully.
    
    Args:
        e: The exception to handle
    """
    logger.error(f'An error occurred: {e}')  # Log the error

class QAExtractor:
    """Main orchestrator for extracting QA audit findings.
    """
    def __init__(self, file_path: str):
        validate_input(file_path)  # Validate the input
        self.file_path = file_path

    def process(self) -> None:
        """Main workflow to extract and process audit findings.
        """
        try:
            raw_data = fetch_data(self.file_path)  # Fetch data from Azure
            records = transform_records(raw_data)  # Transform the data
            save_to_db(records)  # Save to database
        except Exception as e:
            handle_errors(e)  # Handle errors gracefully

if __name__ == '__main__':
    # Example usage
    try:
        extractor = QAExtractor(file_path='path/to/your/audit_report.pdf')
        extractor.process()  # Start the extraction process
    except ValueError as ve:
        logger.error(f'ValueError: {ve}')  # Handle value errors
    except Exception as e:
        logger.error(f'Unhandled error: {e}')  # Handle unanticipated errors