adapter

 avatar
unknown
python
9 months ago
104 kB
21
Indexable
"""
Complete AWS Textract Adapter Training Generator
Supports: Anthropic Claude, OpenAI GPT-4, PyMuPDF for PDF processing
"""

import anthropic
import boto3
import json
import base64
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Optional
import logging
import time
import os
from dotenv import load_dotenv
from datetime import datetime
import concurrent.futures
import threading
from functools import partial
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
)
from validation import LLMOutputValidator, validate_llm_output
import glob

# Load environment variables
load_dotenv()


# Setup logging
def setup_logging(log_dir: str = "logs") -> str:
    """Setup file and console logging"""
    os.makedirs(log_dir, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = os.path.join(log_dir, f"textract_adapter_{timestamp}.log")

    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        handlers=[logging.FileHandler(log_file), logging.StreamHandler()],
    )

    logger = logging.getLogger(__name__)
    logger.info(f"Logging initialized. Log file: {log_file}")
    return log_file


# Initialize logging
LOG_FILE = setup_logging()

try:
    import fitz  # PyMuPDF

    PYMUPDF_AVAILABLE = True
except ImportError:
    PYMUPDF_AVAILABLE = False


# GPU Detection
def check_gpu_availability() -> Dict[str, Any]:
    """Check if GPU is available and return GPU information"""
    gpu_info = {
        "available": False,
        "device": "cpu",
        "device_count": 0,
        "cuda_available": False,
        "mps_available": False,
    }

    try:
        import torch

        gpu_info["cuda_available"] = torch.cuda.is_available()
        gpu_info["mps_available"] = (
            torch.backends.mps.is_available()
            if hasattr(torch.backends, "mps")
            else False
        )

        if gpu_info["cuda_available"]:
            gpu_info["available"] = True
            gpu_info["device"] = "cuda"
            gpu_info["device_count"] = torch.cuda.device_count()
            gpu_info["device_name"] = torch.cuda.get_device_name(0)
        elif gpu_info["mps_available"]:
            gpu_info["available"] = True
            gpu_info["device"] = "mps"
            gpu_info["device_name"] = "Apple Metal Performance Shaders"
        else:
            gpu_info["device"] = "cpu"
            gpu_info["device_name"] = "CPU"

    except ImportError:
        gpu_info["error"] = "PyTorch not installed"

    return gpu_info


# Initialize GPU info
GPU_INFO = check_gpu_availability()

logger = logging.getLogger(__name__)


class TextractAdapterTrainingGenerator:
    """
    Generate training data for AWS Textract adapters using Vision LLMs
    Supports: Claude, OpenAI, PyMuPDF
    """

    def __init__(
        self,
        llm_provider: Optional[str] = None,
        anthropic_api_key: Optional[str] = None,
        openai_api_key: Optional[str] = None,
        aws_access_key_id: Optional[str] = None,
        aws_secret_access_key: Optional[str] = None,
        aws_region: Optional[str] = None,
        aws_profile: Optional[str] = None,
        use_pymupdf: bool = True,
        use_layoutlm: bool = True,  # Enable LayoutLM processing
        openai_response_format: str = "json_object",  # "json_object" or "json_schema"
        max_workers: int = 4,  # Number of threads for concurrent processing
        enable_multithreading: bool = True,  # Enable/disable multithreading
        local_only: bool = False,  # Skip S3 operations for local processing
        enable_validation: bool = True,  # Enable LLM output validation
        validation_threshold: float = 0.7,  # Minimum validation score threshold
    ):
        # Load from environment variables if not provided
        self.llm_provider = (
            llm_provider or os.getenv("LLM_PROVIDER", "anthropic")
        ).lower()
        self.use_pymupdf = use_pymupdf and PYMUPDF_AVAILABLE

        # LayoutLM configuration with GPU checking
        self.use_layoutlm = use_layoutlm
        self.gpu_info = GPU_INFO
        self.layoutlm_available = False
        self.layoutlm_model = None
        self.layoutlm_tokenizer = None

        if self.use_layoutlm:
            self._initialize_layoutlm()

        self.openai_response_format = openai_response_format or os.getenv(
            "OPENAI_RESPONSE_FORMAT", "json_object"
        )

        # Multithreading configuration
        self.max_workers = max_workers
        self.enable_multithreading = enable_multithreading
        self.local_only = local_only
        self._log_lock = threading.Lock()  # Thread-safe logging

        # Validation configuration
        self.enable_validation = enable_validation
        self.validation_threshold = validation_threshold
        self.validator = LLMOutputValidator() if enable_validation else None

        if self.use_pymupdf:
            logger.info("PyMuPDF enabled for PDF processing")
        elif use_pymupdf and not PYMUPDF_AVAILABLE:
            logger.warning(
                "PyMuPDF requested but not installed. Install with: pip install pymupdf"
            )

        if self.llm_provider == "openai":
            logger.info(f"OpenAI response format: {self.openai_response_format}")

        if self.enable_multithreading:
            logger.info(f"Multithreading enabled with {self.max_workers} workers")

        if self.local_only:
            logger.info("Local-only mode: S3 operations disabled")

        if self.enable_validation:
            logger.info(
                f"LLM output validation enabled (threshold: {self.validation_threshold})"
            )
        else:
            logger.info("LLM output validation disabled")

        # Log GPU and LayoutLM status
        if self.gpu_info["available"]:
            logger.info(
                f"GPU detected: {self.gpu_info['device_name']} ({self.gpu_info['device']})"
            )
        else:
            logger.info(f"GPU not available, using: {self.gpu_info['device']}")

        if self.layoutlm_available:
            logger.info("LayoutLM model loaded and ready")
        elif self.use_layoutlm:
            logger.warning("LayoutLM requested but not available")

        # Initialize LLM client
        if self.llm_provider == "anthropic":
            api_key = anthropic_api_key or os.getenv("ANTHROPIC_API_KEY")
            if not api_key:
                raise ValueError(
                    "anthropic_api_key required when using Anthropic. Set ANTHROPIC_API_KEY environment variable or pass anthropic_api_key parameter."
                )
            self.anthropic_client = anthropic.Anthropic(api_key=api_key)
            self.llm_client = None
        elif self.llm_provider == "openai":
            api_key = openai_api_key or os.getenv("OPENAI_API_KEY")
            if not api_key:
                raise ValueError(
                    "openai_api_key required when using OpenAI. Set OPENAI_API_KEY environment variable or pass openai_api_key parameter."
                )
            try:
                from openai import OpenAI

                self.llm_client = OpenAI(api_key=api_key)
                self.anthropic_client = None
            except ImportError:
                raise ImportError(
                    "OpenAI package not installed. Run: pip install openai"
                )
        else:
            raise ValueError(
                f"Unsupported LLM provider: {self.llm_provider}. Use 'anthropic' or 'openai'"
            )

        # AWS clients with environment variable support (skip if local-only mode)
        if self.local_only:
            self.s3 = None
            self.textract = None
            logger.info("AWS clients disabled for local-only mode")
        else:
            aws_access_key = aws_access_key_id or os.getenv("AWS_ACCESS_KEY_ID")
            aws_secret_key = aws_secret_access_key or os.getenv("AWS_SECRET_ACCESS_KEY")
            aws_region = aws_region or os.getenv("AWS_DEFAULT_REGION", "us-east-1")

            if aws_profile or (not aws_access_key and not aws_secret_key):
                # Use AWS profile or default credentials
                session = boto3.Session(profile_name=aws_profile)
                self.s3 = session.client("s3", region_name=aws_region)
                self.textract = session.client("textract", region_name=aws_region)
            else:
                # Use explicit credentials
                if not aws_access_key or not aws_secret_key:
                    raise ValueError(
                        "AWS credentials required. Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or pass aws_access_key_id and aws_secret_access_key parameters."
                    )

                session = boto3.Session(
                    aws_access_key_id=aws_access_key,
                    aws_secret_access_key=aws_secret_key,
                    region_name=aws_region,
                )
                self.s3 = session.client("s3")
                self.textract = session.client("textract")

    def _initialize_layoutlm(self):
        """Initialize LayoutLM model with GPU support"""
        try:
            from transformers import (
                LayoutLMv3ForTokenClassification,
                LayoutLMv3Tokenizer,
            )
            import torch

            # Check if we should use GPU
            device = "cpu"
            if self.gpu_info["available"]:
                device = self.gpu_info["device"]
                logger.info(f"Initializing LayoutLM on {device}")
            else:
                logger.info("GPU not available, initializing LayoutLM on CPU")

            # Load model and tokenizer
            model_name = "microsoft/layoutlmv3-base"
            logger.info(f"Loading LayoutLM model: {model_name}")

            self.layoutlm_tokenizer = LayoutLMv3Tokenizer.from_pretrained(model_name)
            self.layoutlm_model = LayoutLMv3ForTokenClassification.from_pretrained(
                model_name
            )

            # Move model to appropriate device
            if device != "cpu":
                self.layoutlm_model = self.layoutlm_model.to(device)
                logger.info(f"LayoutLM model moved to {device}")

            self.layoutlm_available = True
            logger.info("LayoutLM initialization successful")

        except ImportError as e:
            logger.warning(f"LayoutLM dependencies not available: {e}")
            logger.warning("Install with: pip install transformers torch torchvision")
            self.layoutlm_available = False
        except Exception as e:
            logger.error(f"LayoutLM initialization failed: {e}")
            self.layoutlm_available = False

    def create_document_schema(self, document_type: str) -> Dict[str, Any]:
        """
        Define expected schema for different document types
        Note: LLM will extract ALL fields, not just these - this is just for organization
        """
        schemas = {
            "commercial_insurance": {
                "applicant_info": [
                    "company_name",
                    "mailing_address",
                    "business_phone",
                    "website",
                    "fein",
                    "entity_type",
                    "gl_code",
                    "sic_code",
                    "naics_code",
                ],
                "policy_info": [
                    "policy_number",
                    "effective_date",
                    "expiration_date",
                    "agency_name",
                    "agency_address",
                    "contact_name",
                    "contact_phone",
                    "contact_email",
                    "agency_customer_id",
                    "naic_code",
                    "transaction_type",
                ],
                "premises_info": [
                    "location_number",
                    "location_address",
                    "city",
                    "state",
                    "zip_code",
                    "county",
                    "building_area",
                    "number_employees_fulltime",
                    "number_employees_parttime",
                    "annual_revenues",
                    "construction_type",
                    "year_built",
                    "occupancy_type",
                    "protection_class",
                ],
                "coverage_info": [
                    "building_limit",
                    "bpp_limit",
                    "business_income_limit",
                    "rental_income_limit",
                    "deductible",
                    "causes_of_loss",
                    "coinsurance_percentage",
                    "valuation_type",
                    "earthquake_coverage",
                    "flood_coverage",
                ],
                "prior_carrier": [
                    "year",
                    "carrier_name",
                    "policy_number",
                    "premium",
                    "effective_date",
                    "expiration_date",
                ],
                "loss_history": [
                    "date_of_occurrence",
                    "date_of_claim",
                    "claim_type",
                    "description",
                    "amount_paid",
                    "amount_reserved",
                ],
                "additional_interests": [
                    "interest_type",
                    "name",
                    "address",
                    "loan_number",
                    "location",
                ],
                "contact_info": ["contact_type", "contact_name", "phone", "email"],
            },
            "invoice": {
                "invoice_details": ["invoice_number", "invoice_date", "due_date"],
                "vendor_info": ["vendor_name", "vendor_address", "vendor_contact"],
                "line_items": ["description", "quantity", "unit_price", "total"],
                "totals": ["subtotal", "tax", "total_amount"],
            },
        }

        return schemas.get(document_type, {})

    def generate_llm_prompt(self, document_type: str) -> str:
        """Create document-type-specific prompt for LLM"""
        schema = self.create_document_schema(document_type)

        prompt = f"""You are analyzing a {document_type.replace('_', ' ')} document to extract structured data for training an AWS Textract custom adapter.

Your task is to:
1. Identify and extract ALL key-value pairs and their precise locations
2. For each field, provide normalized bounding box coordinates (0-1 scale)
3. Categorize fields according to the expected schema
4. Identify any tables with their structure and location

Expected Schema Categories:
{json.dumps(schema, indent=2)}

CRITICAL REQUIREMENTS:
- All bbox coordinates MUST be normalized between 0 and 1
- bbox format: {{"x": left, "y": top, "width": width, "height": height}}
- x and y are top-left corner coordinates
- width and height are the dimensions
- Extract EVERY visible field, EVEN IF NOT in the schema above
- For fields not in schema, use category: "other" or create appropriate category
- For multi-line addresses, create separate entries for each component
- For tables, identify header row and data rows separately

Return ONLY valid JSON in this EXACT format:
{{
    "document_type": "{document_type}",
    "fields": [
        {{
            "category": "category_name",
            "field_name": "normalized_field_name",
            "field_value": "extracted_value",
            "bbox": {{"x": 0.0, "y": 0.0, "width": 0.0, "height": 0.0}},
            "confidence": "high|medium|low",
            "page": 1
        }}
    ],
    "tables": [
        {{
            "category": "table_name",
            "bbox": {{"x": 0.0, "y": 0.0, "width": 0.0, "height": 0.0}},
            "rows": 5,
            "columns": 4,
            "page": 1
        }}
    ],
    "document_metadata": {{
        "total_pages": 1,
        "document_quality": "high|medium|low"
    }}
}}

IMPORTANT: Return ONLY the JSON object. Do not include any markdown formatting, code blocks, or additional text."""

        return prompt

    def _thread_safe_log(self, level: str, message: str):
        """Thread-safe logging method"""
        with self._log_lock:
            if level == "info":
                logger.info(message)
            elif level == "warning":
                logger.warning(message)
            elif level == "error":
                logger.error(message)
            elif level == "debug":
                logger.debug(message)

    def validate_llm_annotations(
        self, annotations: Dict[str, Any], document_type: str = "commercial_insurance"
    ) -> Dict[str, Any]:
        """
        Validate LLM-generated annotations using the validation module

        Args:
            annotations: LLM-generated annotations to validate
            document_type: Type of document being processed

        Returns:
            Validation results dictionary
        """
        if not self.enable_validation or not self.validator:
            logger.info("Validation disabled, skipping validation")
            return {"validation_disabled": True}

        logger.info("Starting LLM output validation...")

        try:
            # Update validator document type if needed
            if self.validator.document_type != document_type:
                self.validator.document_type = document_type
                self.validator.validation_rules = (
                    self.validator._load_validation_rules()
                )

            # Run validation
            validation_results = self.validator.validate_annotations(annotations)

            # Check if validation passes threshold
            overall_score = validation_results.get("overall_score", 0)
            validation_passed = overall_score >= self.validation_threshold

            validation_results["validation_passed"] = validation_passed
            validation_results["threshold"] = self.validation_threshold

            # Log validation results
            if validation_passed:
                logger.info(
                    f"āœ… Validation PASSED (score: {overall_score:.2f} >= {self.validation_threshold})"
                )
            else:
                logger.warning(
                    f"āŒ Validation FAILED (score: {overall_score:.2f} < {self.validation_threshold})"
                )

            # Log critical issues
            critical_issues = validation_results.get("critical_issues", [])
            if critical_issues:
                logger.warning(f"🚨 {len(critical_issues)} critical issues found:")
                for issue in critical_issues:
                    logger.warning(f"  - {issue}")

            # Log warnings
            warnings = validation_results.get("warnings", [])
            if warnings:
                logger.info(f"āš ļø  {len(warnings)} warnings found:")
                for warning in warnings:
                    logger.info(f"  - {warning}")

            # Log recommendations
            recommendations = validation_results.get("recommendations", [])
            if recommendations:
                logger.info(f"šŸ’” {len(recommendations)} recommendations:")
                for rec in recommendations:
                    logger.info(f"  - {rec}")

            return validation_results

        except Exception as e:
            logger.error(f"Validation failed with error: {e}")
            return {
                "validation_error": str(e),
                "validation_passed": False,
                "overall_score": 0.0,
            }

    # PyMuPDF Functions
    def get_pdf_info(self, pdf_path: str) -> Dict[str, Any]:
        """Get PDF metadata using PyMuPDF with PDF type detection"""
        if not self.use_pymupdf:
            return {"total_pages": None, "method": "direct"}

        try:
            doc = fitz.open(pdf_path)
            info = {
                "total_pages": len(doc),
                "file_size": Path(pdf_path).stat().st_size,
                "method": "pymupdf",
                "pdf_type": "unknown",
                "pages": [],
            }

            # Analyze PDF type (digital vs scannable)
            total_text_length = 0
            total_images = 0

            for page_num in range(len(doc)):
                page = doc[page_num]
                page_text = page.get_text()
                text_length = len(page_text)
                total_text_length += text_length

                # Count images on page
                image_list = page.get_images()
                total_images += len(image_list)

                info["pages"].append(
                    {
                        "page_number": page_num + 1,
                        "width": page.rect.width,
                        "height": page.rect.height,
                        "text_length": text_length,
                        "image_count": len(image_list),
                        "has_text": text_length > 0,
                    }
                )

            # Determine PDF type
            avg_text_per_page = total_text_length / len(doc) if len(doc) > 0 else 0
            avg_images_per_page = total_images / len(doc) if len(doc) > 0 else 0

            if avg_text_per_page > 100:  # Good amount of text
                info["pdf_type"] = "digital"
                logger.info("šŸ“„ Digital PDF detected - good for direct processing")
            elif avg_images_per_page > 0:  # Has images but little text
                info["pdf_type"] = "scannable"
                logger.info("šŸ“· Scannable PDF detected - may need OCR preprocessing")
            else:
                info["pdf_type"] = "mixed"
                logger.info("šŸ”„ Mixed PDF detected - combination of text and images")

            info["analysis"] = {
                "avg_text_per_page": avg_text_per_page,
                "avg_images_per_page": avg_images_per_page,
                "total_text_length": total_text_length,
                "total_images": total_images,
            }

            doc.close()
            return info

        except Exception as e:
            logger.warning(f"Could not read PDF metadata: {e}")
            return {"total_pages": None, "method": "direct", "pdf_type": "unknown"}

    def pdf_to_images(
        self, pdf_path: str, output_dir: str = None, dpi: int = 150
    ) -> List[str]:
        """Convert PDF pages to images using PyMuPDF"""
        if not self.use_pymupdf:
            raise RuntimeError(
                "PyMuPDF not available. Install with: pip install pymupdf"
            )

        if output_dir is None:
            output_dir = f"{Path(pdf_path).stem}_images"

        Path(output_dir).mkdir(parents=True, exist_ok=True)

        doc = fitz.open(pdf_path)
        image_paths = []

        logger.info(f"Converting {len(doc)} pages to images (DPI: {dpi})...")

        for page_num in range(len(doc)):
            page = doc[page_num]
            zoom = dpi / 72
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)

            image_path = f"{output_dir}/page_{page_num + 1:03d}.png"
            pix.save(image_path)
            image_paths.append(image_path)

            logger.info(f"  Page {page_num + 1}/{len(doc)} -> {image_path}")

        doc.close()
        logger.info(f"Converted {len(image_paths)} pages")

        return image_paths

    def split_pdf(
        self, pdf_path: str, output_dir: str = None, pages_per_split: int = 10
    ) -> List[str]:
        """Split large PDF into smaller chunks"""
        if not self.use_pymupdf:
            raise RuntimeError("PyMuPDF not available")

        if output_dir is None:
            output_dir = f"{Path(pdf_path).stem}_splits"

        Path(output_dir).mkdir(exist_ok=True)

        doc = fitz.open(pdf_path)
        total_pages = len(doc)
        split_files = []

        logger.info(
            f"Splitting {total_pages} pages into chunks of {pages_per_split}..."
        )

        for start_page in range(0, total_pages, pages_per_split):
            end_page = min(start_page + pages_per_split, total_pages)

            output_pdf = fitz.open()
            output_pdf.insert_pdf(doc, from_page=start_page, to_page=end_page - 1)

            split_filename = f"{output_dir}/{Path(pdf_path).stem}_pages_{start_page + 1}_to_{end_page}.pdf"
            output_pdf.save(split_filename)
            output_pdf.close()

            split_files.append(split_filename)
            logger.info(f"  Created: {split_filename}")

        doc.close()
        return split_files

    def optimize_pdf(
        self, pdf_path: str, output_path: str = None, quality: int = 75
    ) -> str:
        """Optimize/compress PDF using PyMuPDF"""
        if not self.use_pymupdf:
            raise RuntimeError("PyMuPDF not available")

        if output_path is None:
            output_path = f"{Path(pdf_path).stem}_optimized.pdf"

        logger.info(f"Optimizing PDF: {pdf_path}")

        doc = fitz.open(pdf_path)
        doc.save(output_path, garbage=4, deflate=True)
        doc.close()

        original_size = Path(pdf_path).stat().st_size
        optimized_size = Path(output_path).stat().st_size
        reduction = (1 - optimized_size / original_size) * 100

        logger.info(
            f"Optimized: {original_size:,} → {optimized_size:,} bytes ({reduction:.1f}% reduction)"
        )

        return output_path

    def auto_annotate_pdf(
        self,
        pdf_path: str,
        output_path: str = None,
        document_type: str = "commercial_insurance",
    ) -> str:
        """Auto-annotate PDF with PyMuPDF, then refine with LLM"""
        if not self.use_pymupdf:
            raise RuntimeError("PyMuPDF not available")

        if output_path is None:
            output_path = f"{Path(pdf_path).stem}_auto_annotated.pdf"

        logger.info(f"Auto-annotating PDF: {pdf_path}")

        doc = fitz.open(pdf_path)
        annotations_added = 0

        # Define common field patterns for different document types
        field_patterns = self._get_field_patterns(document_type)

        for page_num in range(len(doc)):
            page = doc[page_num]
            page_annotations = 0

            # Extract text blocks with positions
            text_blocks = page.get_text("dict")

            for block in text_blocks["blocks"]:
                if "lines" in block:
                    for line in block["lines"]:
                        for span in line["spans"]:
                            text = span["text"].strip()
                            bbox = span["bbox"]

                            # Check if text matches field patterns
                            field_info = self._match_field_pattern(text, field_patterns)

                            if field_info:
                                # Add highlight annotation
                                rect = fitz.Rect(bbox)
                                highlight = page.add_highlight_annot(rect)
                                highlight.set_colors(
                                    stroke=(1, 1, 0)
                                )  # Yellow highlight
                                highlight.set_info(
                                    title=field_info["category"],
                                    content=field_info["field_name"],
                                )

                                # Add text annotation with field info
                                point = fitz.Point(bbox[0], bbox[1] - 10)
                                text_annot = page.add_text_annot(
                                    point,
                                    f"{field_info['field_name']}: {field_info['category']}",
                                )
                                text_annot.set_colors(stroke=(0, 0, 1))  # Blue text

                                page_annotations += 1
                                annotations_added += 1

            logger.info(f"Page {page_num + 1}: Added {page_annotations} annotations")

        doc.save(output_path)
        doc.close()

        logger.info(f"Auto-annotated PDF saved: {output_path}")
        logger.info(f"Total annotations added: {annotations_added}")

        return output_path

    def _get_field_patterns(self, document_type: str) -> Dict[str, List[str]]:
        """Get field patterns for auto-detection"""
        patterns = {
            "commercial_insurance": {
                "company_name": ["company", "corporation", "inc", "llc", "ltd", "co"],
                "policy_number": ["policy", "policy #", "policy number", "pol #"],
                "effective_date": ["effective", "effective date", "inception"],
                "expiration_date": ["expiration", "expires", "exp date"],
                "premium": ["premium", "amount", "cost", "price"],
                "address": ["address", "street", "city", "state", "zip"],
                "phone": ["phone", "telephone", "tel", "call"],
                "email": ["email", "e-mail", "@"],
                "fein": ["fein", "federal id", "tax id", "ein"],
                "naic": ["naic", "naic code"],
                "coverage": ["coverage", "limit", "deductible", "policy limit"],
            },
            "invoice": {
                "invoice_number": ["invoice", "invoice #", "inv #", "bill #"],
                "invoice_date": ["date", "invoice date", "billing date"],
                "due_date": ["due", "due date", "payment due"],
                "vendor": ["vendor", "supplier", "from", "bill to"],
                "amount": ["total", "amount", "sum", "subtotal", "tax"],
                "line_items": ["item", "description", "quantity", "unit price"],
            },
        }
        return patterns.get(document_type, {})

    def _match_field_pattern(
        self, text: str, patterns: Dict[str, List[str]]
    ) -> Optional[Dict[str, str]]:
        """Match text against field patterns"""
        text_lower = text.lower()

        for field_name, keywords in patterns.items():
            for keyword in keywords:
                if keyword in text_lower:
                    return {
                        "field_name": field_name,
                        "category": self._get_category_for_field(field_name),
                        "confidence": "medium",
                    }
        return None

    def _get_category_for_field(self, field_name: str) -> str:
        """Get category for field name"""
        category_map = {
            "company_name": "applicant_info",
            "policy_number": "policy_info",
            "effective_date": "policy_info",
            "expiration_date": "policy_info",
            "premium": "policy_info",
            "address": "applicant_info",
            "phone": "contact_info",
            "email": "contact_info",
            "fein": "applicant_info",
            "naic": "policy_info",
            "coverage": "coverage_info",
        }
        return category_map.get(field_name, "other")

    def hybrid_extract_annotations(
        self,
        pdf_path: str,
        document_type: str = "commercial_insurance",
        save_images: bool = True,
        results_dir: str = None,
        include_layoutlm: bool = True,
    ) -> Dict[str, Any]:
        """Hybrid approach: PyMuPDF auto-annotation + LLM refinement"""
        logger.info(f"Starting hybrid annotation process for {pdf_path}")

        # Create structured directory for this PDF
        pdf_name = Path(pdf_path).stem
        if results_dir:
            pdf_output_dir = Path(results_dir) / "pdf_files" / pdf_name
            pdf_output_dir.mkdir(parents=True, exist_ok=True)

            # Copy original PDF to structured directory
            import shutil

            original_copy = pdf_output_dir / f"{pdf_name}_original.pdf"
            shutil.copy2(pdf_path, original_copy)
            logger.info(f"Saved original PDF: {original_copy}")

        # Step 1: Auto-annotate with PyMuPDF
        annotated_pdf = self.auto_annotate_pdf(pdf_path, document_type=document_type)

        # Save annotated PDF to structured directory
        if results_dir:
            annotated_copy = pdf_output_dir / f"{pdf_name}_annotated.pdf"
            shutil.copy2(annotated_pdf, annotated_copy)
            logger.info(f"Saved annotated PDF: {annotated_copy}")

        # Step 2: Save PyMuPDF images if requested
        if save_images and results_dir:
            self._save_pymupdf_images(pdf_path, pdf_output_dir, annotated_pdf)

        # Step 3: Extract annotations from annotated PDF
        pymupdf_annotations = self._extract_pymupdf_annotations(annotated_pdf)

        # Step 4: Use LLM to refine and categorize
        llm_annotations = self.extract_annotations_with_llm(
            annotated_pdf, document_type, pdf_output_dir
        )

        # Step 5: Optional LayoutLM processing if available and requested
        layoutlm_annotations = {"fields": [], "tables": []}
        if include_layoutlm and self.layoutlm_available and save_images and results_dir:
            logger.info("Running LayoutLM processing on PDF images...")
            try:
                # Get image directory from saved images
                images_dir = pdf_output_dir / "images" / "original"
                if images_dir.exists():
                    image_files = sorted(images_dir.glob("*.png"))
                    logger.info(f"Processing {len(image_files)} images with LayoutLM")

                    for img_path in image_files:
                        page_num = int(
                            img_path.stem.split("_")[-1]
                        )  # Extract page number
                        layoutlm_result = self._extract_with_layoutlm(
                            str(img_path), document_type
                        )

                        # Update page numbers for LayoutLM results
                        for field in layoutlm_result.get("fields", []):
                            field["page"] = page_num
                            layoutlm_annotations["fields"].append(field)

                    logger.info(
                        f"LayoutLM extracted {len(layoutlm_annotations['fields'])} additional fields"
                    )
                else:
                    logger.warning("No images found for LayoutLM processing")
            except Exception as e:
                logger.error(f"LayoutLM processing failed: {e}")

        # Step 6: Combine all annotations
        combined_annotations = self._combine_annotations(
            pymupdf_annotations, llm_annotations, layoutlm_annotations
        )

        # Step 6: Validate combined annotations if enabled
        if self.enable_validation:
            validation_results = self.validate_llm_annotations(
                combined_annotations, document_type
            )
            combined_annotations["validation_results"] = validation_results

            # Save validation report if results directory is provided
            if results_dir:
                self._save_validation_report(pdf_output_dir, validation_results)

        # Step 7: Save all results to structured directory
        if results_dir:
            self._save_structured_results(
                pdf_output_dir,
                pdf_name,
                pymupdf_annotations,
                llm_annotations,
                combined_annotations,
            )

        logger.info(
            f"Hybrid process complete: {len(combined_annotations.get('fields', []))} fields extracted"
        )

        return combined_annotations

    def _save_structured_results(
        self,
        pdf_output_dir: Path,
        pdf_name: str,
        pymupdf_annotations: Dict[str, Any],
        llm_annotations: Dict[str, Any],
        combined_annotations: Dict[str, Any],
    ):
        """Save all results in a structured format"""
        import json
        from datetime import datetime

        # Create subdirectories
        annotations_dir = pdf_output_dir / "annotations"
        metadata_dir = pdf_output_dir / "metadata"
        results_dir = pdf_output_dir / "results"

        annotations_dir.mkdir(exist_ok=True)
        metadata_dir.mkdir(exist_ok=True)
        results_dir.mkdir(exist_ok=True)

        # Save PyMuPDF annotations
        pymupdf_file = annotations_dir / f"{pdf_name}_pymupdf_annotations.json"
        with open(pymupdf_file, "w") as f:
            json.dump(pymupdf_annotations, f, indent=2)
        logger.info(f"Saved PyMuPDF annotations: {pymupdf_file}")

        # Save LLM annotations
        llm_file = annotations_dir / f"{pdf_name}_llm_annotations.json"
        with open(llm_file, "w") as f:
            json.dump(llm_annotations, f, indent=2)
        logger.info(f"Saved LLM annotations: {llm_file}")

        # Save combined annotations (unified ground truth)
        combined_file = annotations_dir / f"{pdf_name}_unified_ground_truth.json"
        with open(combined_file, "w") as f:
            json.dump(combined_annotations, f, indent=2)
        logger.info(f"Saved unified ground truth: {combined_file}")

        # Save metadata
        metadata = {
            "document_name": pdf_name,
            "processing_timestamp": datetime.now().isoformat(),
            "llm_provider": self.llm_provider,
            "multithreading_enabled": self.enable_multithreading,
            "max_workers": self.max_workers,
            "pymupdf_enabled": self.use_pymupdf,
            "total_fields": len(combined_annotations.get("fields", [])),
            "total_tables": len(combined_annotations.get("tables", [])),
            "pymupdf_fields": len(pymupdf_annotations.get("fields", [])),
            "llm_fields": len(llm_annotations.get("fields", [])),
            "document_metadata": combined_annotations.get("document_metadata", {}),
        }

        metadata_file = metadata_dir / f"{pdf_name}_metadata.json"
        with open(metadata_file, "w") as f:
            json.dump(metadata, f, indent=2)
        logger.info(f"Saved metadata: {metadata_file}")

        # Save processing summary
        summary = {
            "document_name": pdf_name,
            "processing_summary": {
                "total_pages": combined_annotations.get("document_metadata", {}).get(
                    "total_pages", 0
                ),
                "total_fields_extracted": len(combined_annotations.get("fields", [])),
                "total_tables_found": len(combined_annotations.get("tables", [])),
                "processing_method": "hybrid_pymupdf_llm",
                "llm_provider": self.llm_provider,
                "multithreading_used": self.enable_multithreading,
            },
            "field_categories": self._categorize_fields(
                combined_annotations.get("fields", [])
            ),
            "files_generated": {
                "original_pdf": f"{pdf_name}_original.pdf",
                "annotated_pdf": f"{pdf_name}_annotated.pdf",
                "pymupdf_annotations": f"{pdf_name}_pymupdf_annotations.json",
                "llm_annotations": f"{pdf_name}_llm_annotations.json",
                "unified_ground_truth": f"{pdf_name}_unified_ground_truth.json",
                "metadata": f"{pdf_name}_metadata.json",
            },
        }

        summary_file = results_dir / f"{pdf_name}_processing_summary.json"
        with open(summary_file, "w") as f:
            json.dump(summary, f, indent=2)
        logger.info(f"Saved processing summary: {summary_file}")

    def _categorize_fields(self, fields: List[Dict[str, Any]]) -> Dict[str, int]:
        """Categorize fields by type for summary"""
        categories = {}
        for field in fields:
            field_type = field.get("field_type", "unknown")
            categories[field_type] = categories.get(field_type, 0) + 1
        return categories

    def _save_pymupdf_images(
        self, original_pdf: str, results_dir: Path, annotated_pdf: str
    ):
        """Save PyMuPDF images (original and annotated) for review"""
        if not self.use_pymupdf:
            return

        doc_name = Path(original_pdf).stem
        images_dir = results_dir / "images"
        images_dir.mkdir(parents=True, exist_ok=True)

        logger.info(f"Saving PyMuPDF images to: {images_dir}")

        # Save original PDF images
        try:
            original_images = self.pdf_to_images(
                original_pdf, str(images_dir / "original"), dpi=150
            )
            logger.info(f"Saved {len(original_images)} original images")
        except Exception as e:
            logger.warning(f"Could not save original images: {e}")

        # Save annotated PDF images
        try:
            annotated_images = self.pdf_to_images(
                annotated_pdf, str(images_dir / "annotated"), dpi=150
            )
            logger.info(f"Saved {len(annotated_images)} annotated images")
        except Exception as e:
            logger.warning(f"Could not save annotated images: {e}")

        # Create comparison HTML file
        self._create_image_comparison_html(images_dir, doc_name)

    def _create_image_comparison_html(self, images_dir: Path, doc_name: str):
        """Create HTML file for side-by-side image comparison"""
        original_dir = images_dir / "original"
        annotated_dir = images_dir / "annotated"

        if not original_dir.exists() or not annotated_dir.exists():
            return

        # Get image files
        original_images = sorted(original_dir.glob("*.png"))
        annotated_images = sorted(annotated_dir.glob("*.png"))

        if not original_images or not annotated_images:
            return

        html_content = f"""
<!DOCTYPE html>
<html>
<head>
    <title>PyMuPDF Image Comparison - {doc_name}</title>
    <style>
        body {{ font-family: Arial, sans-serif; margin: 20px; }}
        .comparison {{ display: flex; margin-bottom: 20px; border: 1px solid #ccc; }}
        .original, .annotated {{ flex: 1; padding: 10px; }}
        .original {{ background-color: #f9f9f9; }}
        .annotated {{ background-color: #fff9e6; }}
        .original h3 {{ color: #333; }}
        .annotated h3 {{ color: #d97706; }}
        img {{ max-width: 100%; height: auto; border: 1px solid #ddd; }}
        .page-info {{ text-align: center; margin: 10px 0; font-weight: bold; }}
    </style>
</head>
<body>
    <h1>PyMuPDF Image Comparison: {doc_name}</h1>
    <p>Side-by-side comparison of original PDF pages vs PyMuPDF annotated pages</p>
"""

        for i, (orig_img, ann_img) in enumerate(zip(original_images, annotated_images)):
            html_content += f"""
    <div class="comparison">
        <div class="original">
            <h3>Original Page {i+1}</h3>
            <div class="page-info">Page {i+1}</div>
            <img src="{orig_img.name}" alt="Original Page {i+1}">
        </div>
        <div class="annotated">
            <h3>PyMuPDF Annotated Page {i+1}</h3>
            <div class="page-info">Page {i+1} (with auto-detected fields)</div>
            <img src="{ann_img.name}" alt="Annotated Page {i+1}">
        </div>
    </div>
"""

        html_content += """
</body>
</html>
"""

        html_file = images_dir / f"{doc_name}_comparison.html"
        with open(html_file, "w") as f:
            f.write(html_content)

        logger.info(f"Created image comparison HTML: {html_file}")

    def _extract_pymupdf_annotations(self, pdf_path: str) -> Dict[str, Any]:
        """Extract annotations from PyMuPDF annotated PDF"""
        if not self.use_pymupdf:
            return {"fields": [], "tables": []}

        doc = fitz.open(pdf_path)
        annotations = {"fields": [], "tables": []}

        for page_num in range(len(doc)):
            page = doc[page_num]
            page_annotations = page.annots()

            for annot in page_annotations:
                if annot.type[0] == 8:  # Highlight annotation
                    rect = annot.rect
                    info = annot.info

                    # Convert to normalized coordinates
                    page_rect = page.rect
                    bbox = {
                        "x": rect.x0 / page_rect.width,
                        "y": rect.y0 / page_rect.height,
                        "width": (rect.x1 - rect.x0) / page_rect.width,
                        "height": (rect.y1 - rect.y0) / page_rect.height,
                    }

                    field = {
                        "category": info.get("title", "unknown"),
                        "field_name": info.get("content", "unknown"),
                        "field_value": "",  # Will be filled by LLM
                        "bbox": bbox,
                        "confidence": "medium",
                        "page": page_num + 1,
                        "source": "pymupdf_auto",
                    }
                    annotations["fields"].append(field)

        doc.close()
        return annotations

    def _combine_annotations(
        self,
        pymupdf_annotations: Dict[str, Any],
        llm_annotations: Dict[str, Any],
        layoutlm_annotations: Dict[str, Any] = None,
    ) -> Dict[str, Any]:
        """Combine PyMuPDF, LLM, and LayoutLM annotations with LLM priority"""
        combined = {
            "document_type": llm_annotations.get(
                "document_type", "commercial_insurance"
            ),
            "fields": [],
            "tables": llm_annotations.get("tables", []),
            "document_metadata": llm_annotations.get("document_metadata", {}),
        }

        pymupdf_fields = pymupdf_annotations.get("fields", [])
        llm_fields = llm_annotations.get("fields", [])
        layoutlm_fields = (
            layoutlm_annotations.get("fields", []) if layoutlm_annotations else []
        )

        # Prioritize LLM annotations - they have higher quality
        logger.info(
            f"Combining annotations: {len(llm_fields)} LLM fields, {len(pymupdf_fields)} PyMuPDF fields, {len(layoutlm_fields)} LayoutLM fields"
        )

        # Start with LLM annotations as primary source
        for llm_field in llm_fields:
            # Create a copy to avoid modifying the original
            llm_field_copy = llm_field.copy()
            # Add source information
            llm_field_copy["source"] = "llm_primary"
            combined["fields"].append(llm_field_copy)

        # Add PyMuPDF fields only if they don't overlap with LLM fields
        # and have meaningful content (not empty field_value)
        for pymupdf_field in pymupdf_fields:
            pymupdf_bbox = pymupdf_field.get("bbox", {})
            pymupdf_value = pymupdf_field.get("field_value", "").strip()

            # Skip empty PyMuPDF fields
            if not pymupdf_value:
                continue

            # Check for overlap with existing LLM fields
            overlaps = False
            for existing_field in combined["fields"]:
                if self._bbox_overlap(pymupdf_bbox, existing_field.get("bbox", {})):
                    overlaps = True
                    break

            if not overlaps:
                pymupdf_field["source"] = "pymupdf_supplement"
                combined["fields"].append(pymupdf_field)

        # Add LayoutLM fields only if they don't overlap with existing fields
        for layoutlm_field in layoutlm_fields:
            layoutlm_bbox = layoutlm_field.get("bbox", {})
            layoutlm_value = layoutlm_field.get("field_value", "").strip()

            # Skip empty LayoutLM fields
            if not layoutlm_value:
                continue

            # Check for overlap with existing fields
            overlaps = False
            for existing_field in combined["fields"]:
                if self._bbox_overlap(layoutlm_bbox, existing_field.get("bbox", {})):
                    overlaps = True
                    break

            if not overlaps:
                layoutlm_field["source"] = "layoutlm_supplement"
                combined["fields"].append(layoutlm_field)

        logger.info(f"Combined result: {len(combined['fields'])} total fields")
        return combined

    def _bbox_overlap(self, bbox1: Dict, bbox2: Dict, threshold: float = 0.3) -> bool:
        """Check if two bounding boxes overlap significantly"""
        if not bbox1 or not bbox2:
            return False

        # Calculate intersection over union (IoU)
        x1 = max(bbox1.get("x", 0), bbox2.get("x", 0))
        y1 = max(bbox1.get("y", 0), bbox2.get("y", 0))
        x2 = min(
            bbox1.get("x", 0) + bbox1.get("width", 0),
            bbox2.get("x", 0) + bbox2.get("width", 0),
        )
        y2 = min(
            bbox1.get("y", 0) + bbox1.get("height", 0),
            bbox2.get("y", 0) + bbox2.get("height", 0),
        )

        if x2 <= x1 or y2 <= y1:
            return False

        intersection = (x2 - x1) * (y2 - y1)
        area1 = bbox1.get("width", 0) * bbox1.get("height", 0)
        area2 = bbox2.get("width", 0) * bbox2.get("height", 0)
        union = area1 + area2 - intersection

        if union == 0:
            return False

        iou = intersection / union
        return iou > threshold

    def _extract_with_layoutlm(
        self, image_path: str, document_type: str = "commercial_insurance"
    ) -> Dict[str, Any]:
        """Extract annotations using Microsoft's LayoutLM model from Hugging Face"""
        if not self.layoutlm_available:
            logger.warning("LayoutLM not available, skipping extraction")
            return {"fields": [], "tables": [], "source": "layoutlm_unavailable"}

        try:
            import torch
            from PIL import Image

            logger.info(f"Processing image with LayoutLM: {image_path}")

            # Load and process image
            image = Image.open(image_path).convert("RGB")

            # Tokenize and process
            inputs = self.layoutlm_tokenizer(
                image, return_tensors="pt", padding=True, truncation=True
            )

            # Move inputs to the same device as model
            device = self.gpu_info["device"] if self.gpu_info["available"] else "cpu"
            inputs = {
                k: v.to(device) if isinstance(v, torch.Tensor) else v
                for k, v in inputs.items()
            }

            # Get predictions
            with torch.no_grad():
                outputs = self.layoutlm_model(**inputs)
                predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
                predicted_labels = torch.argmax(predictions, dim=-1)

            # Convert predictions to annotations
            annotations = {
                "document_type": document_type,
                "fields": [],
                "tables": [],
                "source": "layoutlm",
                "processing_device": device,
            }

            # Process token predictions
            tokens = self.layoutlm_tokenizer.convert_ids_to_tokens(
                inputs["input_ids"][0]
            )
            bboxes = inputs["bbox"][0]

            current_field = None
            for i, (token, bbox, label) in enumerate(
                zip(tokens, bboxes, predicted_labels[0])
            ):
                if token.startswith("##") or token in ["[CLS]", "[SEP]", "[PAD]"]:
                    continue

                # Create field annotation based on label
                if label.item() > 0:  # Non-O label
                    # Normalize bounding box coordinates
                    normalized_bbox = {
                        "x": bbox[0].item() / 1000.0,  # LayoutLM uses 0-1000 scale
                        "y": bbox[1].item() / 1000.0,
                        "width": (bbox[2] - bbox[0]).item() / 1000.0,
                        "height": (bbox[3] - bbox[1]).item() / 1000.0,
                    }

                    field = {
                        "category": "layoutlm_detected",
                        "field_name": f"layoutlm_field_{label.item()}_{i}",
                        "field_value": token,
                        "bbox": normalized_bbox,
                        "confidence": "medium",  # LayoutLM base model has lower confidence
                        "source": "layoutlm",
                        "page": 1,  # Will be updated by caller
                    }
                    annotations["fields"].append(field)

            logger.info(
                f"LayoutLM extracted {len(annotations['fields'])} fields on {device}"
            )
            return annotations

        except Exception as e:
            logger.error(f"LayoutLM extraction failed: {e}")
            return {
                "fields": [],
                "tables": [],
                "source": "layoutlm_error",
                "error": str(e),
            }

    def create_results_folder(self, base_name: str = None) -> str:
        """Create organized results folder structure"""
        if base_name is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            base_name = f"textract_results_{timestamp}"

        results_dir = Path(base_name)
        results_dir.mkdir(exist_ok=True)

        # Create subdirectories
        (results_dir / "annotations").mkdir(exist_ok=True)
        (results_dir / "manifests").mkdir(exist_ok=True)
        (results_dir / "annotated_pdfs").mkdir(exist_ok=True)
        (results_dir / "pymupdf_images").mkdir(exist_ok=True)
        (results_dir / "logs").mkdir(exist_ok=True)
        (results_dir / "raw_llm_output").mkdir(exist_ok=True)

        logger.info(f"Created results folder: {results_dir}")
        return str(results_dir)

    def load_pdfs_from_folder(
        self, folder_path: str, recursive: bool = True
    ) -> List[str]:
        """Load all PDF files from a specified folder"""
        folder_path = Path(folder_path)

        if not folder_path.exists():
            logger.error(f"Folder does not exist: {folder_path}")
            return []

        if not folder_path.is_dir():
            logger.error(f"Path is not a directory: {folder_path}")
            return []

        logger.info(f"Scanning for PDFs in: {folder_path}")

        # Define search pattern
        if recursive:
            pattern = "**/*.pdf"
            logger.info("Searching recursively in subdirectories")
        else:
            pattern = "*.pdf"
            logger.info("Searching only in top-level directory")

        # Find PDF files
        pdf_files = list(folder_path.glob(pattern))

        # Convert to strings and sort
        pdf_paths = sorted([str(pdf) for pdf in pdf_files])

        logger.info(f"Found {len(pdf_paths)} PDF files")
        for pdf_path in pdf_paths:
            file_size = Path(pdf_path).stat().st_size / 1024 / 1024  # MB
            logger.info(f"  - {Path(pdf_path).name} ({file_size:.2f} MB)")

        return pdf_paths

    def save_results_organized(
        self,
        document_paths: List[str],
        s3_bucket: str,
        s3_prefix: str,
        document_type: str = "commercial_insurance",
        use_hybrid_approach: bool = False,
        results_folder: str = None,
        include_layoutlm: bool = False,
    ) -> str:
        """Process documents and save results in organized folders for review"""
        logger.info("Starting organized results processing...")

        # Create results folder
        if results_folder is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            results_folder = f"textract_results_{timestamp}"

        results_dir = self.create_results_folder(results_folder)
        all_manifest_entries = []

        # Process each document
        for doc_path in document_paths:
            try:
                logger.info(f"\n{'='*60}")
                logger.info(f"Processing: {doc_path}")
                logger.info(f"{'='*60}")

                doc_name = Path(doc_path).stem
                doc_filename = Path(doc_path).name

                # Upload to S3
                s3_key = f"{s3_prefix}/documents/{doc_filename}"
                s3_uri = self.upload_to_s3(doc_path, s3_bucket, s3_key)
                logger.info(f"Uploaded to S3: {s3_uri}")

                # Process document
                if use_hybrid_approach:
                    logger.info("Using hybrid PyMuPDF + LLM approach")
                    annotations = self.hybrid_extract_annotations(
                        doc_path,
                        document_type,
                        save_images=True,
                        results_dir=str(results_dir),
                        include_layoutlm=include_layoutlm,
                    )
                else:
                    logger.info("Using LLM-only approach")
                    annotations = self.extract_annotations_with_llm(
                        doc_path, document_type
                    )

                # Save raw LLM annotations
                raw_annotations_file = (
                    Path(results_dir)
                    / "raw_llm_output"
                    / f"{doc_name}_raw_annotations.json"
                )
                with open(raw_annotations_file, "w") as f:
                    json.dump(annotations, f, indent=2)
                logger.info(f"Saved raw annotations: {raw_annotations_file}")

                # Convert to Textract format
                manifest_entries = self.convert_to_textract_format(annotations, s3_uri)

                # Save individual manifest
                individual_manifest = (
                    Path(results_dir) / "manifests" / f"{doc_name}_manifest.jsonl"
                )
                with open(individual_manifest, "w") as f:
                    for entry in manifest_entries:
                        f.write(json.dumps(entry) + "\n")
                logger.info(f"Saved individual manifest: {individual_manifest}")

                # Save processed annotations (cleaned up)
                processed_annotations = {
                    "document_name": doc_name,
                    "document_type": document_type,
                    "processing_method": (
                        "hybrid" if use_hybrid_approach else "llm_only"
                    ),
                    "total_fields": len(annotations.get("fields", [])),
                    "total_tables": len(annotations.get("tables", [])),
                    "fields": annotations.get("fields", []),
                    "tables": annotations.get("tables", []),
                    "metadata": annotations.get("document_metadata", {}),
                    "s3_uri": s3_uri,
                    "processing_timestamp": datetime.now().isoformat(),
                }

                processed_file = (
                    Path(results_dir) / "annotations" / f"{doc_name}_processed.json"
                )
                with open(processed_file, "w") as f:
                    json.dump(processed_annotations, f, indent=2)
                logger.info(f"Saved processed annotations: {processed_file}")

                all_manifest_entries.extend(manifest_entries)

                # Copy annotated PDF if hybrid approach was used
                if use_hybrid_approach:
                    annotated_pdf = f"{doc_name}_auto_annotated.pdf"
                    if os.path.exists(annotated_pdf):
                        import shutil

                        dest_pdf = Path(results_dir) / "annotated_pdfs" / annotated_pdf
                        shutil.copy2(annotated_pdf, dest_pdf)
                        logger.info(f"Saved annotated PDF: {dest_pdf}")

            except Exception as e:
                logger.error(f"Error processing {doc_path}: {e}")
                import traceback

                traceback.print_exc()
                continue

        # Save combined manifest
        combined_manifest = (
            Path(results_dir) / "manifests" / "combined_training_manifest.jsonl"
        )
        with open(combined_manifest, "w") as f:
            for entry in all_manifest_entries:
                f.write(json.dumps(entry) + "\n")

        # Create summary report
        summary = {
            "processing_timestamp": datetime.now().isoformat(),
            "total_documents": len(document_paths),
            "total_manifest_entries": len(all_manifest_entries),
            "processing_method": "hybrid" if use_hybrid_approach else "llm_only",
            "document_type": document_type,
            "s3_bucket": s3_bucket,
            "s3_prefix": s3_prefix,
            "results_folder": results_dir,
            "files_created": {
                "raw_annotations": len(
                    list(Path(results_dir).glob("raw_llm_output/*.json"))
                ),
                "processed_annotations": len(
                    list(Path(results_dir).glob("annotations/*.json"))
                ),
                "individual_manifests": len(
                    list(Path(results_dir).glob("manifests/*_manifest.jsonl"))
                ),
                "annotated_pdfs": len(
                    list(Path(results_dir).glob("annotated_pdfs/*.pdf"))
                ),
            },
        }

        summary_file = Path(results_dir) / "processing_summary.json"
        with open(summary_file, "w") as f:
            json.dump(summary, f, indent=2)

        logger.info(f"\n{'='*60}")
        logger.info(f"RESULTS SAVED FOR REVIEW")
        logger.info(f"{'='*60}")
        logger.info(f"Results folder: {results_dir}")
        logger.info(f"Total manifest entries: {len(all_manifest_entries)}")
        logger.info(f"Summary: {summary_file}")
        logger.info(f"Combined manifest: {combined_manifest}")
        logger.info(f"{'='*60}\n")

        return str(results_dir)

    # LLM Processing Functions
    def extract_annotations_with_llm(
        self,
        file_path: str,
        document_type: str = "commercial_insurance",
        results_dir: Path = None,
    ) -> Dict[str, Any]:
        """Use Vision LLM to extract annotations from document"""
        logger.info(f"Processing {file_path} with {self.llm_provider.upper()}...")

        pdf_info = self.get_pdf_info(file_path)
        total_pages = pdf_info.get("total_pages")
        if total_pages:
            logger.info(f"Document has {total_pages} pages")

        file_ext = Path(file_path).suffix.lower()

        # For OpenAI with PDF, convert to images
        if self.llm_provider == "openai" and file_ext == ".pdf":
            if self.use_pymupdf:
                logger.info("Converting PDF to images for OpenAI...")
                # Use structured directory for temporary images
                if results_dir:
                    temp_images_dir = results_dir / "temp_images"
                else:
                    temp_images_dir = None
                image_paths = self.pdf_to_images(
                    file_path,
                    str(temp_images_dir) if temp_images_dir else None,
                    dpi=150,
                )

                all_annotations = {
                    "document_type": document_type,
                    "fields": [],
                    "tables": [],
                    "document_metadata": {
                        "total_pages": len(image_paths),
                        "document_quality": "high",
                        "processing_method": "pymupdf_conversion",
                    },
                }

                if self.enable_multithreading and len(image_paths) > 1:
                    # Use multithreading for multiple pages
                    self._thread_safe_log(
                        "info",
                        f"Processing {len(image_paths)} pages with multithreading...",
                    )
                    page_annotations_list = self._process_pages_multithreaded(
                        image_paths, document_type, results_dir
                    )

                    for idx, page_annotations in enumerate(page_annotations_list):
                        for field in page_annotations.get("fields", []):
                            field["page"] = idx + 1
                            all_annotations["fields"].append(field)

                        for table in page_annotations.get("tables", []):
                            table["page"] = idx + 1
                            all_annotations["tables"].append(table)
                else:
                    # Sequential processing
                    for idx, img_path in enumerate(image_paths):
                        self._thread_safe_log(
                            "info", f"Processing page {idx + 1}/{len(image_paths)}..."
                        )
                        page_annotations = self._process_single_file(
                            img_path, document_type, page_number=idx + 1
                        )

                        for field in page_annotations.get("fields", []):
                            field["page"] = idx + 1
                            all_annotations["fields"].append(field)

                        for table in page_annotations.get("tables", []):
                            table["page"] = idx + 1
                            all_annotations["tables"].append(table)

                # Cleanup temporary images
                import shutil

                if temp_images_dir and temp_images_dir.exists():
                    shutil.rmtree(temp_images_dir)
                    logger.info("Cleaned up temporary images")
                elif image_paths:
                    # Fallback to old cleanup method
                    shutil.rmtree(Path(image_paths[0]).parent)
                    logger.info("Cleaned up temporary images")

                return all_annotations
            else:
                logger.warning(
                    "PyMuPDF not available. OpenAI may not process PDF correctly."
                )

        # Process the file
        annotations = self._process_single_file(file_path, document_type)

        # Validate annotations if enabled
        if self.enable_validation:
            validation_results = self.validate_llm_annotations(
                annotations, document_type
            )
            annotations["validation_results"] = validation_results

            # Save validation report if results directory is provided
            if results_dir:
                self._save_validation_report(results_dir, validation_results)

        return annotations

    def _process_single_file(
        self, file_path: str, document_type: str, page_number: Optional[int] = None
    ) -> Dict[str, Any]:
        """Process a single file"""
        with open(file_path, "rb") as f:
            file_data = f.read()

        file_b64 = base64.standard_b64encode(file_data).decode("utf-8")

        file_ext = Path(file_path).suffix.lower()
        media_type_map = {
            ".pdf": "application/pdf",
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
        }
        media_type = media_type_map.get(file_ext, "application/pdf")

        prompt = self.generate_llm_prompt(document_type)
        if page_number:
            prompt = f"[PROCESSING PAGE {page_number}]\n\n{prompt}"

        if self.llm_provider == "anthropic":
            annotations = self._call_anthropic(file_b64, media_type, prompt)
        elif self.llm_provider == "openai":
            if self.openai_response_format == "json_schema":
                annotations = self._call_openai_with_schema(
                    file_b64, media_type, prompt, document_type
                )
            else:
                annotations = self._call_openai(file_b64, media_type, prompt)

        return annotations

    def _process_pages_multithreaded(
        self, image_paths: List[str], document_type: str, results_dir: Path = None
    ) -> List[Dict[str, Any]]:
        """Process multiple pages concurrently using ThreadPoolExecutor"""
        page_annotations_list = [None] * len(image_paths)

        def process_single_page(idx_and_path):
            idx, img_path = idx_and_path
            try:
                self._thread_safe_log(
                    "info", f"Processing page {idx + 1}/{len(image_paths)}..."
                )
                page_annotations = self._process_single_file(
                    img_path, document_type, page_number=idx + 1
                )

                # Save individual page results immediately if results_dir is provided
                if results_dir:
                    self._save_individual_page_result(
                        results_dir, idx + 1, page_annotations
                    )

                return idx, page_annotations
            except Exception as e:
                self._thread_safe_log("error", f"Error processing page {idx + 1}: {e}")
                return idx, {"fields": [], "tables": [], "error": str(e)}

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=self.max_workers
        ) as executor:
            # Submit all tasks
            future_to_idx = {
                executor.submit(process_single_page, (idx, img_path)): idx
                for idx, img_path in enumerate(image_paths)
            }

            # Collect results as they complete
            for future in concurrent.futures.as_completed(future_to_idx):
                idx, page_annotations = future.result()
                page_annotations_list[idx] = page_annotations

        return page_annotations_list

    def _save_individual_page_result(
        self, results_dir: Path, page_number: int, page_annotations: Dict[str, Any]
    ):
        """Save individual page LLM results immediately"""
        import json

        # Create raw_llm_output directory if it doesn't exist
        raw_output_dir = results_dir / "raw_llm_output"
        raw_output_dir.mkdir(exist_ok=True)

        # Save individual page result
        page_file = raw_output_dir / f"page_{page_number:03d}_llm_result.json"
        with open(page_file, "w") as f:
            json.dump(page_annotations, f, indent=2)

        self._thread_safe_log(
            "info", f"Saved page {page_number} LLM result: {page_file}"
        )

    def _save_validation_report(
        self, results_dir: Path, validation_results: Dict[str, Any]
    ):
        """Save validation report to results directory"""
        try:
            # Create validation reports directory
            validation_dir = results_dir / "validation_reports"
            validation_dir.mkdir(exist_ok=True)

            # Save detailed validation report
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            report_file = validation_dir / f"validation_report_{timestamp}.json"

            with open(report_file, "w") as f:
                json.dump(validation_results, f, indent=2)

            self._thread_safe_log("info", f"Saved validation report: {report_file}")

            # Save summary report
            summary = {
                "timestamp": validation_results.get("timestamp"),
                "overall_score": validation_results.get("overall_score", 0),
                "validation_passed": validation_results.get("validation_passed", False),
                "critical_issues_count": len(
                    validation_results.get("critical_issues", [])
                ),
                "warnings_count": len(validation_results.get("warnings", [])),
                "recommendations_count": len(
                    validation_results.get("recommendations", [])
                ),
                "validation_summary": validation_results.get("validation_summary", {}),
            }

            summary_file = validation_dir / f"validation_summary_{timestamp}.json"
            with open(summary_file, "w") as f:
                json.dump(summary, f, indent=2)

            self._thread_safe_log("info", f"Saved validation summary: {summary_file}")

        except Exception as e:
            self._thread_safe_log("error", f"Failed to save validation report: {e}")

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type((Exception,)),
        reraise=True,
    )
    def _call_anthropic(
        self, file_b64: str, media_type: str, prompt: str
    ) -> Dict[str, Any]:
        """Call Anthropic Claude API with retry logic"""
        message = self.anthropic_client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=16000,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": (
                                "document"
                                if media_type == "application/pdf"
                                else "image"
                            ),
                            "source": {
                                "type": "base64",
                                "media_type": media_type,
                                "data": file_b64,
                            },
                        },
                        {"type": "text", "text": prompt},
                    ],
                }
            ],
        )

        response_text = message.content[0].text
        return self._parse_llm_response(response_text)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type((Exception,)),
        reraise=True,
    )
    def _call_openai(
        self, file_b64: str, media_type: str, prompt: str
    ) -> Dict[str, Any]:
        """Call OpenAI GPT-4 Vision API with structured JSON output and retry logic"""
        response = self.llm_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{media_type};base64,{file_b64}",
                                "detail": "high",
                            },
                        },
                    ],
                }
            ],
            max_tokens=16000,
            response_format={"type": "json_object"},
        )

        response_text = response.choices[0].message.content

        # With structured output, we can directly parse JSON
        try:
            return json.loads(response_text)
        except json.JSONDecodeError as e:
            logger.warning(f"Failed to parse structured JSON response: {e}")
            # Fallback to the old parsing method
            return self._parse_llm_response(response_text)

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type((Exception,)),
        reraise=True,
    )
    def _call_openai_with_schema(
        self, file_b64: str, media_type: str, prompt: str, document_type: str
    ) -> Dict[str, Any]:
        """Call OpenAI GPT-4 Vision API with JSON schema for even more structured output and retry logic"""

        # Define the JSON schema for the expected output
        json_schema = {
            "type": "object",
            "properties": {
                "document_type": {"type": "string"},
                "fields": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "category": {"type": "string"},
                            "field_name": {"type": "string"},
                            "field_value": {"type": "string"},
                            "bbox": {
                                "type": "object",
                                "properties": {
                                    "x": {"type": "number", "minimum": 0, "maximum": 1},
                                    "y": {"type": "number", "minimum": 0, "maximum": 1},
                                    "width": {
                                        "type": "number",
                                        "minimum": 0,
                                        "maximum": 1,
                                    },
                                    "height": {
                                        "type": "number",
                                        "minimum": 0,
                                        "maximum": 1,
                                    },
                                },
                                "required": ["x", "y", "width", "height"],
                            },
                            "confidence": {
                                "type": "string",
                                "enum": ["high", "medium", "low"],
                            },
                            "page": {"type": "integer", "minimum": 1},
                        },
                        "required": [
                            "category",
                            "field_name",
                            "field_value",
                            "bbox",
                            "confidence",
                            "page",
                        ],
                    },
                },
                "tables": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "category": {"type": "string"},
                            "bbox": {
                                "type": "object",
                                "properties": {
                                    "x": {"type": "number", "minimum": 0, "maximum": 1},
                                    "y": {"type": "number", "minimum": 0, "maximum": 1},
                                    "width": {
                                        "type": "number",
                                        "minimum": 0,
                                        "maximum": 1,
                                    },
                                    "height": {
                                        "type": "number",
                                        "minimum": 0,
                                        "maximum": 1,
                                    },
                                },
                                "required": ["x", "y", "width", "height"],
                            },
                            "rows": {"type": "integer", "minimum": 1},
                            "columns": {"type": "integer", "minimum": 1},
                            "page": {"type": "integer", "minimum": 1},
                        },
                        "required": ["category", "bbox", "rows", "columns", "page"],
                    },
                },
                "document_metadata": {
                    "type": "object",
                    "properties": {
                        "total_pages": {"type": "integer", "minimum": 1},
                        "document_quality": {
                            "type": "string",
                            "enum": ["high", "medium", "low"],
                        },
                    },
                    "required": ["total_pages", "document_quality"],
                },
            },
            "required": ["document_type", "fields", "tables", "document_metadata"],
        }

        response = self.llm_client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{media_type};base64,{file_b64}",
                                "detail": "high",
                            },
                        },
                    ],
                }
            ],
            max_tokens=16000,
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": f"{document_type}_extraction",
                    "schema": json_schema,
                },
            },
        )

        response_text = response.choices[0].message.content

        # With JSON schema, the response should be perfectly structured
        try:
            return json.loads(response_text)
        except json.JSONDecodeError as e:
            logger.warning(f"Failed to parse JSON schema response: {e}")
            # Fallback to the old parsing method
            return self._parse_llm_response(response_text)

    def _parse_llm_response(self, response_text: str) -> Dict[str, Any]:
        """Parse JSON response from LLM"""
        response_text = response_text.strip()
        if "```json" in response_text:
            response_text = response_text.split("```json")[1].split("```")[0].strip()
        elif "```" in response_text:
            response_text = response_text.split("```")[1].split("```")[0].strip()

        return json.loads(response_text)

    # Textract Format Conversion
    def convert_to_textract_format(
        self, llm_annotations: Dict[str, Any], source_s3_uri: str
    ) -> List[Dict[str, Any]]:
        """Convert LLM annotations to Textract adapter training format"""
        total_pages = llm_annotations.get("document_metadata", {}).get("total_pages", 1)
        manifest_entries = []

        for page_num in range(1, total_pages + 1):
            page_fields = [
                f
                for f in llm_annotations.get("fields", [])
                if f.get("page", 1) == page_num
            ]
            page_tables = [
                t
                for t in llm_annotations.get("tables", [])
                if t.get("page", 1) == page_num
            ]

            if not page_fields and not page_tables:
                continue

            textract_format = {
                "source-ref": source_s3_uri,
                "page": page_num,
                "textract-custom-adapter": {
                    "version": "2023-07-15",
                    "document-type": llm_annotations.get(
                        "document_type", "commercial_insurance"
                    ),
                    "key-value-pairs": [],
                    "tables": [],
                    "blocks": [],
                },
                "textract-custom-adapter-metadata": {
                    "confidence-score": 0.95,
                    "human-validated": False,
                    "generation-method": f"llm-{self.llm_provider}",
                    "page": page_num,
                    "total-pages": total_pages,
                    "fields-on-page": len(page_fields),
                    "tables-on-page": len(page_tables),
                },
            }

            for idx, field in enumerate(page_fields):
                key_block_id = f"key_{page_num}_{idx}"
                value_block_id = f"value_{page_num}_{idx}"

                kv_pair = {
                    "key": {
                        "text": field["field_name"],
                        "geometry": {"boundingBox": field["bbox"]},
                        "id": key_block_id,
                    },
                    "value": {
                        "text": field["field_value"],
                        "geometry": {"boundingBox": field["bbox"]},
                        "id": value_block_id,
                    },
                    "category": field.get("category", "general"),
                    "confidence": field.get("confidence", "high"),
                }
                textract_format["textract-custom-adapter"]["key-value-pairs"].append(
                    kv_pair
                )

                textract_format["textract-custom-adapter"]["blocks"].extend(
                    [
                        {
                            "id": key_block_id,
                            "blockType": "KEY_VALUE_SET",
                            "text": field["field_name"],
                            "geometry": {"boundingBox": field["bbox"]},
                            "entityTypes": ["KEY"],
                        },
                        {
                            "id": value_block_id,
                            "blockType": "KEY_VALUE_SET",
                            "text": field["field_value"],
                            "geometry": {"boundingBox": field["bbox"]},
                            "entityTypes": ["VALUE"],
                        },
                    ]
                )

            for idx, table in enumerate(page_tables):
                table_id = f"table_{page_num}_{idx}"
                table_entry = {
                    "id": table_id,
                    "geometry": {"boundingBox": table["bbox"]},
                    "rows": table["rows"],
                    "columns": table["columns"],
                    "category": table.get("category", "table"),
                }
                textract_format["textract-custom-adapter"]["tables"].append(table_entry)

                textract_format["textract-custom-adapter"]["blocks"].append(
                    {
                        "id": table_id,
                        "blockType": "TABLE",
                        "geometry": {"boundingBox": table["bbox"]},
                        "rowCount": table["rows"],
                        "columnCount": table["columns"],
                    }
                )

            manifest_entries.append(textract_format)

        return manifest_entries

    def upload_to_s3(
        self, local_path: str, bucket: str, key: str, force_upload: bool = False
    ) -> str:
        """Upload file to S3, checking if it already exists"""
        s3_uri = f"s3://{bucket}/{key}"

        # Check if file already exists in S3
        if not force_upload:
            try:
                self.s3.head_object(Bucket=bucket, Key=key)
                logger.info(f"File already exists in S3: {s3_uri}")
                return s3_uri
            except self.s3.exceptions.NoSuchKey:
                pass  # File doesn't exist, proceed with upload
            except Exception as e:
                logger.warning(f"Could not check if file exists in S3: {e}")

        logger.info(f"Uploading {local_path} to {s3_uri}")

        # Get file size for progress tracking
        file_size = os.path.getsize(local_path)
        logger.info(f"File size: {file_size / 1024 / 1024:.2f} MB")

        # Upload with progress tracking for large files
        if file_size > 100 * 1024 * 1024:  # 100MB
            logger.info("Large file detected, using multipart upload...")
            self._upload_large_file(local_path, bucket, key)
        else:
            self.s3.upload_file(local_path, bucket, key)

        logger.info(f"Upload completed: {s3_uri}")
        return s3_uri

    def _upload_large_file(self, local_path: str, bucket: str, key: str):
        """Upload large files using multipart upload"""
        from boto3.s3.transfer import TransferConfig

        config = TransferConfig(
            multipart_threshold=1024 * 25,  # 25MB
            max_concurrency=10,
            multipart_chunksize=1024 * 25,
            use_threads=True,
        )

        self.s3.upload_file(local_path, bucket, key, Config=config)

    def process_documents_batch(
        self,
        document_paths: List[str],
        s3_bucket: str,
        s3_prefix: str,
        document_type: str = "commercial_insurance",
        output_manifest_path: str = "training_manifest.jsonl",
        use_hybrid_approach: bool = False,
    ) -> str:
        """Process multiple documents and create training manifest"""
        all_manifest_entries = []

        for doc_path in document_paths:
            try:
                logger.info(f"\n{'='*60}")
                logger.info(f"Processing: {doc_path}")
                logger.info(f"{'='*60}")

                doc_filename = Path(doc_path).name
                s3_key = f"{s3_prefix}/documents/{doc_filename}"
                s3_uri = self.upload_to_s3(doc_path, s3_bucket, s3_key)

                # Choose processing method
                if use_hybrid_approach:
                    logger.info("Using hybrid PyMuPDF + LLM approach")
                    llm_annotations = self.hybrid_extract_annotations(
                        doc_path, document_type
                    )
                else:
                    logger.info("Using LLM-only approach")
                    llm_annotations = self.extract_annotations_with_llm(
                        doc_path, document_type
                    )

                total_pages = llm_annotations.get("document_metadata", {}).get(
                    "total_pages", 1
                )
                logger.info(f"Document has {total_pages} pages")

                manifest_entries = self.convert_to_textract_format(
                    llm_annotations, s3_uri
                )

                logger.info(f"Generated {len(manifest_entries)} manifest entries")
                all_manifest_entries.extend(manifest_entries)

                annotation_filename = f"{Path(doc_path).stem}_annotations.json"
                with open(annotation_filename, "w") as f:
                    json.dump(llm_annotations, f, indent=2)
                logger.info(f"Saved raw annotations to {annotation_filename}")

                groundtruth_filename = f"{Path(doc_path).stem}_groundtruth.jsonl"
                with open(groundtruth_filename, "w") as f:
                    for entry in manifest_entries:
                        f.write(json.dumps(entry) + "\n")
                logger.info(f"Saved ground truth to {groundtruth_filename}")

            except Exception as e:
                logger.error(f"Error processing {doc_path}: {e}")
                import traceback

                traceback.print_exc()
                continue

        with open(output_manifest_path, "w") as f:
            for entry in all_manifest_entries:
                f.write(json.dumps(entry) + "\n")

        logger.info(f"\n{'='*60}")
        logger.info(f"TRAINING DATA GENERATION COMPLETE")
        logger.info(f"{'='*60}")
        logger.info(f"Total manifest entries: {len(all_manifest_entries)}")
        logger.info(f"Manifest saved to: {output_manifest_path}")
        logger.info(f"{'='*60}\n")

        manifest_s3_key = f"{s3_prefix}/manifests/{Path(output_manifest_path).name}"
        manifest_s3_uri = self.upload_to_s3(
            output_manifest_path, s3_bucket, manifest_s3_key
        )

        logger.info(f"Manifest uploaded to: {manifest_s3_uri}")

        return manifest_s3_uri

    def create_textract_adapter(
        self,
        adapter_name: str,
        manifest_s3_uri: str,
        output_s3_bucket: str,
        output_s3_prefix: str,
        description: str = "Auto-generated adapter from LLM annotations",
    ) -> Dict[str, str]:
        """Create Textract adapter"""
        logger.info(f"Creating Textract adapter: {adapter_name}")

        response = self.textract.create_adapter(
            AdapterName=adapter_name,
            Description=description,
            FeatureTypes=["QUERIES", "FORMS", "TABLES"],
            AutoUpdate={"UpdateLevel": "ENABLED"},
        )

        adapter_id = response["AdapterId"]
        logger.info(f"Created adapter: {adapter_id}")

        manifest_parts = manifest_s3_uri.replace("s3://", "").split("/", 1)
        manifest_bucket = manifest_parts[0]
        manifest_key = manifest_parts[1]

        version_response = self.textract.create_adapter_version(
            AdapterId=adapter_id,
            DatasetConfig={
                "ManifestS3Object": {"Bucket": manifest_bucket, "Name": manifest_key}
            },
            OutputConfig={"S3Bucket": output_s3_bucket, "S3Prefix": output_s3_prefix},
        )

        logger.info(f"Created adapter version: {version_response['AdapterVersion']}")

        return {
            "adapter_id": adapter_id,
            "adapter_version": version_response["AdapterVersion"],
            "status": "TRAINING",
        }

    def test_adapter_extraction(
        self,
        adapter_id: str,
        adapter_version: str,
        test_document_path: str,
        test_s3_bucket: str,
        test_s3_key: str,
    ) -> Dict[str, Any]:
        """Test the trained adapter"""
        logger.info(f"Testing adapter on: {test_document_path}")

        logger.info(f"Uploading test document to s3://{test_s3_bucket}/{test_s3_key}")
        self.s3.upload_file(test_document_path, test_s3_bucket, test_s3_key)

        logger.info("Checking adapter status...")
        max_retries = 60
        retry_interval = 60

        for attempt in range(max_retries):
            adapter_status = self.textract.get_adapter_version(
                AdapterId=adapter_id, AdapterVersion=adapter_version
            )

            status = adapter_status["Status"]
            logger.info(f"Adapter status: {status}")

            if status == "ACTIVE":
                logger.info("Adapter is ready!")
                break
            elif status in ["CREATION_FAILED", "CREATION_ERROR"]:
                raise Exception(f"Adapter creation failed: {status}")
            else:
                logger.info(
                    f"Adapter training... (attempt {attempt + 1}/{max_retries})"
                )
                time.sleep(retry_interval)
        else:
            raise TimeoutError("Adapter training timed out")

        logger.info("Starting Textract analysis with custom adapter...")

        response = self.textract.start_document_analysis(
            DocumentLocation={
                "S3Object": {"Bucket": test_s3_bucket, "Name": test_s3_key}
            },
            FeatureTypes=["FORMS", "TABLES"],
            AdaptersConfig={
                "Adapters": [{"AdapterId": adapter_id, "Version": adapter_version}]
            },
        )

        job_id = response["JobId"]
        logger.info(f"Analysis job started: {job_id}")

        while True:
            response = self.textract.get_document_analysis(JobId=job_id)
            status = response["JobStatus"]

            if status == "SUCCEEDED":
                logger.info("Analysis completed!")
                break
            elif status == "FAILED":
                raise Exception("Textract analysis failed")
            else:
                logger.info(f"Analysis in progress... Status: {status}")
                time.sleep(5)

        results = self._parse_textract_results(response)

        results_file = f"{Path(test_document_path).stem}_adapter_test_results.json"
        with open(results_file, "w") as f:
            json.dump(results, f, indent=2)

        logger.info(f"Test results saved to: {results_file}")

        return results

    def _parse_textract_results(
        self, textract_response: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Parse Textract response"""
        results = {"key_value_pairs": [], "tables": [], "raw_text": []}

        blocks = textract_response.get("Blocks", [])

        key_map = {}
        value_map = {}
        block_map = {block["Id"]: block for block in blocks}

        for block in blocks:
            if block["BlockType"] == "KEY_VALUE_SET":
                if "KEY" in block.get("EntityTypes", []):
                    key_map[block["Id"]] = block
                elif "VALUE" in block.get("EntityTypes", []):
                    value_map[block["Id"]] = block
            elif block["BlockType"] == "LINE":
                results["raw_text"].append(block.get("Text", ""))
            elif block["BlockType"] == "TABLE":
                results["tables"].append(
                    {
                        "id": block["Id"],
                        "rows": block.get("RowCount", 0),
                        "columns": block.get("ColumnCount", 0),
                        "confidence": block.get("Confidence", 0),
                    }
                )

        for key_id, key_block in key_map.items():
            key_text = self._get_text_from_block(key_block, block_map)
            value_text = ""

            if "Relationships" in key_block:
                for relationship in key_block["Relationships"]:
                    if relationship["Type"] == "VALUE":
                        for value_id in relationship["Ids"]:
                            if value_id in value_map:
                                value_text = self._get_text_from_block(
                                    value_map[value_id], block_map
                                )

            results["key_value_pairs"].append(
                {
                    "key": key_text,
                    "value": value_text,
                    "confidence": key_block.get("Confidence", 0),
                }
            )

        return results

    def _get_text_from_block(self, block: Dict, block_map: Dict) -> str:
        """Extract text from block"""
        text = ""
        if "Relationships" in block:
            for relationship in block["Relationships"]:
                if relationship["Type"] == "CHILD":
                    for child_id in relationship["Ids"]:
                        child_block = block_map.get(child_id, {})
                        if child_block.get("BlockType") == "WORD":
                            text += child_block.get("Text", "") + " "
        return text.strip()


def main():
    # Load configuration from environment variables
    llm_provider = os.getenv("LLM_PROVIDER", "anthropic")
    s3_bucket = os.getenv("S3_BUCKET", "your-textract-bucket-name")
    s3_prefix = os.getenv("S3_PREFIX", "textract_documents")
    aws_profile = os.getenv("AWS_PROFILE")

    # Initialize generator with environment variables
    generator = TextractAdapterTrainingGenerator(
        llm_provider=llm_provider,
        aws_profile=aws_profile,
        use_pymupdf=True,
        use_layoutlm=True,  # Enable LayoutLM if GPU is available
        enable_validation=True,  # Enable validation by default
        validation_threshold=0.7,  # 70% threshold
    )

    # Ask user for PDF source
    print("\nPDF Source Options:")
    print("1. Current directory (default)")
    print("2. Specific folder")
    source_choice = input("Choose PDF source (1 or 2): ").strip()

    if source_choice == "2":
        folder_path = input("Enter folder path: ").strip()
        if not folder_path:
            folder_path = "."

        print(f"Scanning folder: {folder_path}")
        document_paths = generator.load_pdfs_from_folder(folder_path)

        if not document_paths:
            print(f"No PDF files found in {folder_path}")
            return 1
    else:
        # Default: current directory
        document_paths = [
            "25-25 Middendorfs PROP - ACORD.pdf",
            "25-26 ACORD PROP - Ed White.pdf",
        ]

        # Filter to only existing files
        existing_docs = [doc for doc in document_paths if os.path.exists(doc)]
        if not existing_docs:
            print("No PDF documents found in current directory.")
            print("Please ensure the following files exist:")
            for doc in document_paths:
                print(f"  - {doc}")
            return 1

        document_paths = existing_docs

    print("=" * 70)
    print(f"TEXTRACT ADAPTER TRAINING - Using {llm_provider.upper()}")
    if generator.use_pymupdf:
        print("PyMuPDF: ENABLED")
    if generator.enable_validation:
        print(f"Validation: ENABLED (threshold: {generator.validation_threshold})")
    print("=" * 70)

    # Ask user about processing method
    print("\nProcessing Methods:")
    print("1. LLM-only (default)")
    print("2. Hybrid PyMuPDF + LLM (experimental)")
    if generator.layoutlm_available:
        print("3. Hybrid PyMuPDF + LLM + LayoutLM (GPU required)")
    else:
        print("3. Hybrid PyMuPDF + LLM + LayoutLM (āš ļø  GPU not available)")

    method_choice = input("Choose processing method (1, 2, or 3): ").strip()

    use_hybrid = method_choice in ["2", "3"]
    include_layoutlm = method_choice == "3"

    if method_choice == "1":
        print("šŸ¤– Using LLM-only approach")
    elif method_choice == "2":
        print("šŸš€ Using HYBRID approach: PyMuPDF auto-annotation + LLM refinement")
    elif method_choice == "3":
        if generator.layoutlm_available:
            print("šŸš€ Using ENHANCED HYBRID approach: PyMuPDF + LLM + LayoutLM")
            print(
                f"   GPU: {generator.gpu_info['device_name']} ({generator.gpu_info['device']})"
            )
        else:
            print("āš ļø  LayoutLM not available (no GPU), falling back to hybrid approach")
            include_layoutlm = False
    else:
        print("Invalid choice, using LLM-only approach")
        use_hybrid = False
        include_layoutlm = False

    # Get PDF info
    for doc_path in document_paths:
        if Path(doc_path).suffix.lower() == ".pdf":
            info = generator.get_pdf_info(doc_path)
            print(f"\n{Path(doc_path).name}:")
            print(f"  Pages: {info.get('total_pages', 'unknown')}")
            print(f"  Size: {info.get('file_size', 0) / 1024 / 1024:.2f} MB")
            print(f"  Type: {info.get('pdf_type', 'unknown').upper()}")

            analysis = info.get("analysis", {})
            if analysis:
                print(
                    f"  Avg text per page: {analysis.get('avg_text_per_page', 0):.0f} chars"
                )
                print(
                    f"  Avg images per page: {analysis.get('avg_images_per_page', 0):.1f}"
                )

                if info.get("pdf_type") == "scannable":
                    print(
                        "  āš ļø  Scannable PDF - consider OCR preprocessing for better results"
                    )
                elif info.get("pdf_type") == "digital":
                    print("  āœ… Digital PDF - optimal for processing")

    print("\n" + "=" * 70)
    print(f"Processing {len(document_paths)} documents...")
    print("=" * 70)

    try:
        # Use organized results approach (no adapter creation)
        results_folder = generator.save_results_organized(
            document_paths=document_paths,
            s3_bucket=s3_bucket,
            s3_prefix=s3_prefix,
            document_type="commercial_insurance",
            use_hybrid_approach=use_hybrid,
            include_layoutlm=include_layoutlm,
        )

        print("\n" + "=" * 70)
        print("PROCESSING COMPLETE - RESULTS SAVED FOR REVIEW")
        print("=" * 70)
        print(f"Results folder: {results_folder}")
        print(f"Log file: {LOG_FILE}")

        print(f"\nšŸ”§ To create adapter later, use:")
        print(f"   python create_adapter.py --results-folder {results_folder}")

        print(f"\nšŸ“Š Logs saved to: {LOG_FILE}")

    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback

        traceback.print_exc()
        return 1

    return 0


if __name__ == "__main__":
    exit(main())
Editor is loading...
Leave a Comment