""" Background worker for processing document generation jobs using batched Claude API. Runs as RQ worker process. """ import asyncio import io import json import os import pathlib import tempfile import time import traceback import zipfile import shutil import base64 import math from typing import Dict, Any, List, Callable from datetime import datetime # Add worker startup logging from .config import settings from .supabase_client import supabase_client from .google_drive import GoogleDriveClient from .utils import ( download_image_to_base64, create_token_mapping_json, download_seed_images, build_prompt, extract_html_documents_from_response, extract_ground_truth, extract_css_from_html, increase_handwriting_font_size, unmark_visual_elements, render_html_to_pdf, preprocess_html_for_pdf, extract_bboxes_from_rendered_pdf, extract_all_bboxes_from_pdf, extract_raw_annotations_from_geometries, process_stage3_complete, process_stage4_ocr, process_stage5_complete, validate_html_structure, validate_pdf, validate_bboxes, retry_on_network_error ) from docgenie.generation.pipeline_01.claude_batching import ClaudeBatchedClient from docgenie import ENV # ==================== Worker Logging Configuration ==================== # Read from environment variable, default to False for cleaner logs VERBOSE_LOGGING = os.getenv('WORKER_VERBOSE_LOGGING', 'false').lower() in ('true', '1', 'yes') def log_verbose(message: str): """Log message only if verbose logging is enabled""" if VERBOSE_LOGGING: print(message) # ==================== Startup Validation ==================== def validate_worker_config(): """Validate worker configuration at startup""" print("=" * 60) print("🔧 Worker Configuration Check") print("=" * 60) # Check Anthropic API if settings.ANTHROPIC_API_KEY: print("✓ ANTHROPIC_API_KEY: Set") else: print("✗ ANTHROPIC_API_KEY: NOT SET (REQUIRED)") # Check Supabase if settings.SUPABASE_URL and settings.SUPABASE_KEY: print(f"✓ SUPABASE: {settings.SUPABASE_URL[:30]}...") else: print("✗ SUPABASE: NOT SET (REQUIRED)") # Check Google OAuth (optional, for token refresh) if settings.GOOGLE_CLIENT_ID and settings.GOOGLE_CLIENT_SECRET: print(f"✓ GOOGLE_CLIENT_ID: {settings.GOOGLE_CLIENT_ID[:20]}...") print("✓ GOOGLE_CLIENT_SECRET: Set") print(" → Token auto-refresh: ENABLED") else: print("⚠ GOOGLE_CLIENT_ID/SECRET: Not set") print(" → Token auto-refresh: DISABLED") print(" → Users must provide fresh access tokens that don't expire during processing") print("=" * 60) # Run validation on module import validate_worker_config() async def process_document_generation_job_async(request_id: str, request_data: Dict[str, Any]): """ Async background job function - processes document generation using batched Claude API. This function: 1. Creates Claude batch with single message (generates N documents) 2. Polls batch until completion 3. Processes all documents (PDFs, handwriting, etc.) 4. Uploads ZIP to user's Google Drive 5. Updates Supabase with results Args: request_id: Document request UUID from Supabase request_data: Request parameters dict containing: - user_id: int - seed_images: List[str] (URLs) - prompt_params: Dict (language, doc_type, num_solutions, etc.) Raises: Exception: Any error during processing (logged to Supabase) """ user_id = request_data['user_id'] google_drive_token = request_data.get('google_drive_token') if google_drive_token == "string": google_drive_token = None google_drive_refresh_token = request_data.get('google_drive_refresh_token') if google_drive_refresh_token == "string": google_drive_refresh_token = None seed_image_urls = request_data['seed_images'] prompt_params = request_data['prompt_params'] # Step 0: Clean up any old generated documents for this request (clean retry) log_verbose(f"[Job {request_id}] Cleaning up old results for request {request_id}...") try: supabase_client.delete_generated_documents(request_id) except Exception as cleanup_err: print(f"[Job {request_id}] ⚠ Cleanup of old records failed: {cleanup_err}") # Validate Google Drive credentials configuration if google_drive_refresh_token: if not settings.GOOGLE_CLIENT_ID or not settings.GOOGLE_CLIENT_SECRET: print(f"[Job {request_id}] ⚠️ WARNING: refresh_token provided but GOOGLE_CLIENT_ID/SECRET not configured") print(f"[Job {request_id}] Token auto-refresh will fail. Ensure access token remains valid.") # Create temporary directories for this job with tempfile.TemporaryDirectory() as tmp_dir: tmp_path = pathlib.Path(tmp_dir) batch_dir = tmp_path / "batches" message_dir = tmp_path / "messages" batch_dir.mkdir(exist_ok=True) message_dir.mkdir(exist_ok=True) # Initialize DatasetExporter for organized structure from .dataset_exporter import DatasetExporter exporter = DatasetExporter(tmp_path, dataset_name="docgenie_documents") try: # ==================== Update Status: Downloading ==================== retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "processing")) print(f"[Job {request_id}] Status: processing (fetching seed images)") # ==================== Step 1: Download Seed Images ==================== log_verbose(f"[Job {request_id}] Downloading {len(seed_image_urls)} seed images...") seed_images_base64 = retry_on_network_error(lambda: download_seed_images(seed_image_urls)) log_verbose(f"[Job {request_id}] Downloaded {len(seed_images_base64)} images") # ==================== Step 2: Build Prompts (Chunked) ==================== prompt_template_path = ENV.PROMPT_TEMPLATES_DIR / "ClaudeRefined12" / "seed-based-json.txt" if not prompt_template_path.exists(): raise FileNotFoundError(f"Prompt template not found: {prompt_template_path}") num_solutions = prompt_params.get('num_solutions', 1) chunk_size = settings.BATCH_PROMPT_CHUNK_SIZE num_prompts = math.ceil(num_solutions / chunk_size) prompts = [] images_base64_list = [] image_docids_list = [] for i in range(num_prompts): # Calculate how many solutions for this specific prompt current_num_solutions = min(chunk_size, num_solutions - (i * chunk_size)) p = build_prompt( language=prompt_params.get('language', 'English'), doc_type=prompt_params.get('doc_type', 'business and administrative'), gt_type=prompt_params.get('gt_type', 'Questions and answers'), gt_format=prompt_params.get('gt_format', '{"question": "answer"}'), num_solutions=current_num_solutions, num_seed_images=len(seed_images_base64), prompt_template_path=prompt_template_path, enable_visual_elements=prompt_params.get('enable_visual_elements', False), visual_element_types=prompt_params.get('visual_element_types', []) ) prompts.append(p) images_base64_list.append(seed_images_base64) image_docids_list.append(["seed"] * len(seed_images_base64)) log_verbose(f"[Job {request_id}] Created {num_prompts} prompts (chunk size: {chunk_size})") # ==================== Step 3: Create Claude Batch ==================== log_verbose(f"[Job {request_id}] Creating Claude batch with {num_prompts} messages...") client = ClaudeBatchedClient(api_key=settings.ANTHROPIC_API_KEY) # Send batch with multiple messages (one per chunk) client.send_batch( model=settings.CLAUDE_MODEL, prompts=prompts, images_base64=images_base64_list, image_docids=image_docids_list, batch_data_directory=batch_dir, max_tokens=16384 ) print(f"[Job {request_id}] ⏳ Batch created with {num_prompts} tasks, awaiting processing...") # ==================== Step 4: Poll Batch Until Complete ==================== client.await_batches( batch_data_directory=batch_dir, message_data_directory=message_dir, sleep_seconds_between_batch=2, sleep_seconds_iteration=settings.BATCH_POLL_INTERVAL ) print(f"[Job {request_id}] ✓ Batch complete") # ==================== Step 5: Read Batch Results ==================== message_files = list(message_dir.glob("*.json")) if not message_files: raise RuntimeError("No message results found after batch completion") html_documents = [] for msg_file in message_files: try: message_data = json.loads(msg_file.read_text()) if message_data.get('result_type') == 'succeeded': llm_response = message_data['response'] docs = extract_html_documents_from_response(llm_response) html_documents.extend(docs) # Extract token usage and track cost (Research Parity) from .utils import calculate_message_cost i_tokens = message_data.get('usage_input_tokens', 0) o_tokens = message_data.get('usage_output_tokens', 0) c_create = message_data.get('cache_creation_input_tokens', 0) c_read = message_data.get('cache_read_input_tokens', 0) cost = calculate_message_cost( model=settings.CLAUDE_MODEL, input_tokens=i_tokens, output_tokens=o_tokens, cache_creation_input_tokens=c_create, cache_read_input_tokens=c_read ) exporter.add_cost(cost, i_tokens, o_tokens, c_create, c_read) log_verbose(f" ✓ Extracted {len(docs)} documents from task {msg_file.stem} (Cost: ${cost:.4f})") else: error_msg = message_data.get('error', 'Unknown error') print(f"[Job {request_id}] ⚠ Task {msg_file.stem} failed: {error_msg}") except Exception as e: print(f"[Job {request_id}] ⚠ Error reading message result {msg_file.name}: {e}") if not html_documents: raise RuntimeError("No valid HTML documents found in any batch results") print(f"[Job {request_id}] ✓ Combined total of {len(html_documents)} documents from all tasks") # ==================== Update Status: Generating ==================== retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "generating")) print(f"[Job {request_id}] Status: generating (processing documents)") # ==================== Step 7: Download Assets from Supabase ==================== assets_temp_dir = None try: assets_path = f"{user_id}/{request_id}/assets" files = retry_on_network_error(lambda: supabase_client.list_files("doc_storage", assets_path)) # Filter out directories (ensure files is a list) asset_files = [f for f in files if f and f.get('id') is not None] if files else [] if asset_files: assets_temp_dir = pathlib.Path(tempfile.mkdtemp()) print(f"[Job {request_id}] Found {len(asset_files)} assets in storage, downloading...") for file_info in asset_files: file_name = file_info['name'] try: file_content = retry_on_network_error(lambda: supabase_client.download_file("doc_storage", f"{assets_path}/{file_name}")) with open(assets_temp_dir / file_name, 'wb') as f: f.write(file_content) log_verbose(f" ✓ Downloaded {file_name}") except Exception as download_err: print(f" ⚠ Failed to download {file_name}: {download_err}") else: log_verbose(f"[Job {request_id}] No assets found in {assets_path}") except Exception as e: print(f"[Job {request_id}] ⚠ Asset check/download failed: {e}") # ==================== Step 8: Process Each Document ==================== pdf_files = [] metadata = [] for idx, html in enumerate(html_documents): try: doc_id = f"document_{idx + 1}" log_verbose(f"[Job {request_id}] Processing document {idx + 1}/{len(html_documents)}") # Initialize original_pdf_path original_pdf_path = None # Validate HTML is_valid, error_msg = validate_html_structure(html) if not is_valid: print(f"[Job {request_id}] Document {idx + 1} HTML validation failed: {error_msg}") continue # Extract ground truth and CSS gt, html_clean = extract_ground_truth(html) css, _ = extract_css_from_html(html_clean) # Render to PDF pdf_path = tmp_path / f"{doc_id}.pdf" pdf_path, width_mm, height_mm, geometries = await render_html_to_pdf( html=html_clean, output_pdf_path=pdf_path ) # Track original PDF original_pdf_path = pdf_path # Validate PDF is_valid, error_msg = validate_pdf(pdf_path) if not is_valid: print(f"[Job {request_id}] Document {idx + 1} PDF validation failed: {error_msg}") continue # Extract bounding boxes bboxes_raw = extract_bboxes_from_rendered_pdf(pdf_path) # Validate bboxes is_valid, error_msg = validate_bboxes(bboxes_raw, min_bbox_count=1) if not is_valid: print(f"[Job {request_id}] Document {idx + 1} BBox validation warning: {error_msg}") log_verbose(f"[Job {request_id}] Document {idx + 1}: Extracted {len(bboxes_raw)} bboxes") # Process Stage 3 (Handwriting & Visual Elements) if enabled final_image_b64 = None handwriting_regions = [] visual_elements = [] handwriting_images = {} visual_element_images = {} ocr_results = None pdf_with_handwriting_path = None pdf_final_path = None if prompt_params.get('enable_handwriting') or prompt_params.get('enable_visual_elements'): # Update status: Handwriting if prompt_params.get('enable_handwriting'): retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "handwriting")) log_verbose(f"[Job {request_id}] Status: handwriting (generating handwritten text)") log_verbose(f"[Job {request_id}] Document {idx + 1}: Processing handwriting/visual elements...") try: final_image_b64, handwriting_regions, visual_elements, handwriting_images, visual_element_images, pdf_with_handwriting_path, pdf_final_path = await process_stage3_complete( pdf_path=pdf_path, geometries=geometries, ground_truth=gt, bboxes_raw=bboxes_raw, page_width_mm=width_mm, page_height_mm=height_mm, enable_handwriting=prompt_params.get('enable_handwriting', False), handwriting_ratio=prompt_params.get('handwriting_ratio', 0.3), handwriting_apply_ink_filter=prompt_params.get('handwriting_apply_ink_filter', True), handwriting_enable_enhancements=prompt_params.get('handwriting_enable_enhancements', False), handwriting_num_inference_steps=prompt_params.get('handwriting_num_inference_steps', 1000), handwriting_writer_ids=prompt_params.get('handwriting_writer_ids', [404, 347, 156, 253, 354, 166, 320]), enable_visual_elements=prompt_params.get('enable_visual_elements', False), visual_element_types=prompt_params.get('visual_element_types', []), seed=prompt_params.get('seed'), assets_dir=assets_temp_dir, barcode_number=prompt_params.get('barcode_number') ) # Use final PDF if both modifications applied, otherwise use handwriting PDF if pdf_final_path and pdf_final_path.exists(): pdf_path = pdf_final_path elif pdf_with_handwriting_path and pdf_with_handwriting_path.exists(): pdf_path = pdf_with_handwriting_path log_verbose(f"[Job {request_id}] Document {idx + 1}: {len(handwriting_regions)} handwriting, {len(visual_elements)} visual elements") except Exception as e: print(f"[Job {request_id}] Document {idx + 1}: Stage 3 failed: {str(e)}") # Process Stage 4/5 (OCR) if needed if prompt_params.get('enable_ocr'): # Update status: OCR retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "ocr")) log_verbose(f"[Job {request_id}] Status: ocr (running OCR on documents)") log_verbose(f"[Job {request_id}] Document {idx + 1}: Processing OCR...") try: stage4_image, ocr_results = await process_stage4_ocr( pdf_path=pdf_path, enable_ocr=True, dpi=settings.OCR_DPI ) if ocr_results: log_verbose(f"[Job {request_id}] Document {idx + 1}: OCR complete - {len(ocr_results.get('words', []))} words") except Exception as e: print(f"[Job {request_id}] Document {idx + 1}: OCR failed: {str(e)}") # Process Stage 5 (Dataset packaging) if needed stage5_results = {} if any([ prompt_params.get('enable_bbox_normalization'), prompt_params.get('enable_gt_verification'), prompt_params.get('enable_analysis'), prompt_params.get('enable_debug_visualization') ]): # Update status: Validation (if GT verification enabled) if prompt_params.get('enable_gt_verification'): retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "validation")) log_verbose(f"[Job {request_id}] Status: validation (validating ground truth)") log_verbose(f"[Job {request_id}] Document {idx + 1}: Processing dataset packaging...") try: stage5_results = await process_stage5_complete( document_id=doc_id, pdf_path=pdf_path, image_base64=final_image_b64, ocr_results=ocr_results, ground_truth=gt, has_handwriting=prompt_params.get('enable_handwriting', False), has_visual_elements=prompt_params.get('enable_visual_elements', False), layout_elements=visual_elements, enable_bbox_normalization=prompt_params.get('enable_bbox_normalization', False), enable_gt_verification=prompt_params.get('enable_gt_verification', False), enable_analysis=prompt_params.get('enable_analysis', False), enable_debug_visualization=prompt_params.get('enable_debug_visualization', False) ) except Exception as e: print(f"[Job {request_id}] Document {idx + 1}: Stage 5 failed: {str(e)}") # Track PDFs for metadata if original_pdf_path and pdf_path != original_pdf_path: pdf_files.append(original_pdf_path) pdf_files.append(pdf_path) else: pdf_files.append(pdf_path) # Extract bbox_pdf (word + char) from original PDF (ground truth positions) log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Extracting bbox_pdf (word + char level) from original PDF...") try: bboxes_pdf = extract_all_bboxes_from_pdf(original_pdf_path if original_pdf_path else pdf_path) bbox_pdf_word = bboxes_pdf.get('word', []) bbox_pdf_char = bboxes_pdf.get('char', []) log_verbose(f"[Job {request_id}] Document {idx + 1}: ✓ Extracted {len(bbox_pdf_word)} word bboxes, {len(bbox_pdf_char)} char bboxes from PDF") except Exception as e: print(f"[Job {request_id}] Document {idx + 1}: ⚠ bbox_pdf extraction failed: {e}") bbox_pdf_word = bboxes_raw # Fallback to raw bboxes bbox_pdf_char = [] # Extract raw_annotations (layout boxes before normalization) raw_annotations = None if geometries: log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Extracting raw_annotations from geometries...") try: raw_annotations = extract_raw_annotations_from_geometries(geometries) log_verbose(f"[Job {request_id}] Document {idx + 1}: ✓ Extracted {len(raw_annotations)} layout annotations") except Exception as e: print(f"[Job {request_id}] Document {idx + 1}: ⚠ raw_annotations extraction failed: {e}") # Decode final image to bytes final_image_bytes = None if final_image_b64: import base64 final_image_bytes = base64.b64decode(final_image_b64) # Decode debug visualization debug_viz_bytes = None if stage5_results.get('debug_visualization'): import base64 debug_viz_dict = stage5_results['debug_visualization'] if debug_viz_dict and 'bbox_overlay_base64' in debug_viz_dict: debug_viz_b64 = debug_viz_dict['bbox_overlay_base64'] debug_viz_bytes = base64.b64decode(debug_viz_b64) # Prepare token mapping if tokens exist output_detail = prompt_params.get('output_detail', 'minimal') token_mapping_data = None if output_detail in ["dataset", "complete"]: token_mapping_data = create_token_mapping_json( handwriting_regions, handwriting_images, visual_elements, visual_element_images ) log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Output detail '{output_detail}': Prepared {len(handwriting_images)} handwriting tokens, {len(visual_element_images)} visual elements") # Extract bbox_final_word and bbox_final_segment (from OCR or PDF) bbox_final_word = None bbox_final_segment = None if ocr_results and ocr_results.get('words'): # Use OCR results as final bboxes bbox_final_word = ocr_results.get('words', []) bbox_final_segment = ocr_results.get('lines', []) else: # Fallback to PDF bboxes if no OCR bbox_final_word = bbox_pdf_word bbox_final_segment = [] # No line-level data without OCR # Read PDF bytes for exporter pdf_initial_bytes = original_pdf_path.read_bytes() # Read modified PDFs if they exist pdf_with_handwriting_bytes = None pdf_final_bytes = None pdf_with_visual_elements_bytes = None if pdf_with_handwriting_path and pdf_with_handwriting_path.exists(): pdf_with_handwriting_bytes = pdf_with_handwriting_path.read_bytes() if pdf_final_path and pdf_final_path.exists(): pdf_final_bytes = pdf_final_path.read_bytes() # Special case: if only visual elements (no handwriting), pdf_final is actually pdf_with_visual_elements if pdf_final_bytes and not pdf_with_handwriting_bytes: pdf_with_visual_elements_bytes = pdf_final_bytes pdf_final_bytes = None # Add document to exporter log_verbose(f"[Job {request_id}] Document {idx + 1}: 📦 Adding document to dataset exporter...") exporter.add_document( document_id=doc_id, html=html_clean, css=css, pdf_initial=pdf_initial_bytes, pdf_with_handwriting=pdf_with_handwriting_bytes, pdf_with_visual_elements=pdf_with_visual_elements_bytes, pdf_final=pdf_final_bytes, final_image=final_image_bytes, ground_truth=gt, raw_annotations=raw_annotations, bboxes_pdf_word=bbox_pdf_word, bboxes_pdf_char=bbox_pdf_char, bboxes_final_word=bbox_final_word, bboxes_final_segment=bbox_final_segment, bboxes_normalized_word=stage5_results.get('normalized_bboxes_word'), bboxes_normalized_segment=stage5_results.get('normalized_bboxes_segment'), gt_verification=stage5_results.get('gt_verification'), token_mapping=token_mapping_data, handwriting_regions=handwriting_regions, handwriting_images=handwriting_images, visual_elements=visual_elements, visual_element_images=visual_element_images, layout_elements=visual_elements, geometries=geometries, ocr_results=ocr_results, analysis_stats=stage5_results.get('analysis_stats'), debug_visualization=debug_viz_bytes ) log_verbose(f"[Job {request_id}] Document {idx + 1}: ✓ Document {doc_id} added to dataset") # Store comprehensive metadata (matching /generate/pdf format) metadata.append({ "document_id": doc_id, "filename": f"{doc_id}.pdf", "bboxes": bboxes_raw, "ground_truth": gt, "geometries": geometries, "page_width_mm": width_mm, "page_height_mm": height_mm, "handwriting_regions": handwriting_regions, "visual_elements": visual_elements, "has_stage3_image": final_image_b64 is not None, "ocr_results": ocr_results, # Stage 5 results "normalized_bboxes_word": stage5_results.get('normalized_bboxes_word'), "normalized_bboxes_segment": stage5_results.get('normalized_bboxes_segment'), "gt_verification": stage5_results.get('gt_verification'), "analysis_stats": stage5_results.get('analysis_stats'), "debug_visualization_available": stage5_results.get('debug_visualization') is not None }) except Exception as e: print(f"[Job {request_id}] Error processing document {idx + 1}: {str(e)}") traceback.print_exc() continue if not pdf_files: raise RuntimeError("Failed to process any documents") log_verbose(f"[Job {request_id}] Processed {len(pdf_files)} PDF files") # ==================== Step 8: Finalize Dataset & Create ZIP ==================== log_verbose(f"[Job {request_id}] 📦 Finalizing dataset export...") exporter.finalize( request_id=request_id, user_id=user_id, prompt_params=prompt_params, api_mode="async" ) log_verbose(f"[Job {request_id}] ✓ Dataset structure finalized at {exporter.base_path}") # ==================== Update Status: Zipping ==================== retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "zipping")) print(f"[Job {request_id}] Status: zipping (creating ZIP archive)") # Create ZIP from organized dataset log_verbose(f"[Job {request_id}] 📦 Creating ZIP archive from dataset...") zip_path = tmp_path / f"docgenie_{request_id}.zip" with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_file: # Add all files from exporter.base_path for file_path in exporter.base_path.rglob('*'): if file_path.is_file(): arcname = file_path.relative_to(exporter.base_path.parent) zip_file.write(file_path, arcname) zip_size_mb = zip_path.stat().st_size / (1024 * 1024) log_verbose(f"[Job {request_id}] ✓ ZIP created: {zip_size_mb:.2f} MB") # ==================== Update Status: Uploading ==================== retry_on_network_error(lambda: supabase_client.update_request_status(request_id, "uploading")) print(f"[Job {request_id}] Status: uploading (uploading to Google Drive)") # ==================== Step 9: Upload to Google Drive ==================== print(f"[Job {request_id}] ⬆️ Uploading to Google Drive...") google_drive_url = None gdrive_failed = False gdrive_skipped = False # Check if Google Drive token provided if not google_drive_token or google_drive_token == "string": print(f"[Job {request_id}] No valid Google Drive token provided. Skipping Google Drive upload.") gdrive_skipped = True else: try: drive_client = GoogleDriveClient( access_token=google_drive_token, refresh_token=google_drive_refresh_token ) google_drive_url = drive_client.upload_file( file_path=zip_path, filename=f"docgenie_{request_id}.zip", folder_name=settings.GOOGLE_DRIVE_FOLDER_NAME ) print(f"[Job {request_id}] ✓ Uploaded to Google Drive: {google_drive_url}") except Exception as e: print(f"[Job {request_id}] Google Drive upload failed: {str(e)}") gdrive_failed = True # Do not raise an error, just continue so we can still save to Supabase # ==================== Step 10: Store Results in Supabase ==================== log_verbose(f"[Job {request_id}] Saving results to Supabase...") log_verbose(f"[Job {request_id}] URL: {google_drive_url}") # Upload ZIP to Supabase zip_url = None try: zip_storage_path = f"{user_id}/{request_id}/generated/docgenie_{request_id}.zip" retry_on_network_error(lambda: supabase_client.upload_to_storage("doc_storage", zip_storage_path, zip_path.read_bytes(), "application/zip")) zip_url = supabase_client.get_public_url("doc_storage", zip_storage_path) print(f"[Job {request_id}] ✓ Uploaded ZIP to Supabase: {zip_url}") except Exception as e: print(f"[Job {request_id}] ⚠ Supabase ZIP upload failed: {e}") # ==================== Step 11: Upload Individual Documents to Supabase ==================== print(f"[Job {request_id}] Uploading individual documents to Supabase...") for idx, doc_data in enumerate(metadata): doc_id = doc_data["document_id"] try: # Determine paths (matching sync endpoint structure) doc_storage_path = f"{user_id}/{request_id}/generated/{idx}_doc.pdf" gt_storage_path = f"{user_id}/{request_id}/generated/{idx}_gt.json" src_storage_path = f"{user_id}/{request_id}/generated/{idx}_src.html" bbox_storage_path = f"{user_id}/{request_id}/generated/{idx}_bbox.json" # Find files on disk doc_path = exporter.pdf_final_dir / f"{doc_id}.pdf" if not doc_path.exists(): doc_path = exporter.pdf_initial_dir / f"{doc_id}.pdf" gt_path = exporter.gt_dir / f"{doc_id}.json" src_path = exporter.html_dir / f"{doc_id}.html" bbox_path = exporter.bbox_pdf_word_dir / f"{doc_id}.json" # Step 10: Upload Individual Files and Create Record # We wrap each upload in a retry, and use a nested try-except for the whole group # to ensure that if one document fails, we still try to process others. try: # Upload PDF (Critical) if doc_path.exists(): retry_on_network_error(lambda: supabase_client.upload_to_storage("doc_storage", doc_storage_path, doc_path.read_bytes(), "application/pdf")) # Upload Ground Truth (Important) if gt_path.exists(): retry_on_network_error(lambda: supabase_client.upload_to_storage("doc_storage", gt_storage_path, gt_path.read_bytes(), "application/json")) # Upload HTML Source (Optional) if src_path.exists(): retry_on_network_error(lambda: supabase_client.upload_to_storage("doc_storage", src_storage_path, src_path.read_bytes(), "text/html")) # Upload Bounding Boxes (Optional) if bbox_path.exists(): retry_on_network_error(lambda: supabase_client.upload_to_storage("doc_storage", bbox_storage_path, bbox_path.read_bytes(), "application/json")) # Upload visual element images if available if doc_data.get("visual_elements") and doc_data.get("visual_element_images"): for ve_id, img_b64 in doc_data["visual_element_images"].items(): ve_storage_path = f"{user_id}/{request_id}/generated/{idx}_ve_{ve_id}.png" try: img_bytes = base64.b64decode(img_b64) retry_on_network_error(lambda: supabase_client.upload_to_storage("doc_storage", ve_storage_path, img_bytes, "image/png")) except Exception as ve_err: print(f" ⚠ Failed to upload visual element {ve_id}: {ve_err}") except Exception as upload_err: # Log error but try to create the DB record anyway with what we have print(f" ⚠ Some file uploads failed for document {idx+1}: {upload_err}") # Create record in database (Always try this) try: log_verbose(f" 📦 Creating DB record for document {idx+1} (index {idx})...") record_id = retry_on_network_error(lambda: supabase_client.create_generated_document( request_id=request_id, file_url=supabase_client.get_public_url("doc_storage", doc_storage_path), model_version=settings.LLM_MODEL, doc_index=idx, doc_storage_path=doc_storage_path, gt_storage_path=gt_storage_path, html_storage_path=src_storage_path, bbox_storage_path=bbox_storage_path )) print(f" ✓ Processed document {idx+1} and created DB record {record_id}") except Exception as db_err: print(f" ❌ Failed to create DB record for document {idx+1}: {db_err}") except Exception as doc_err: print(f" ❌ Unexpected error processing document {idx+1}: {doc_err}") # ==================== Step 11: Finalize Request Status ==================== if gdrive_skipped: final_status = "completed_no_gdrive" elif gdrive_failed: final_status = "completed_gdrive_failed" else: final_status = "completed" retry_on_network_error(lambda: supabase_client.update_request_status( request_id=request_id, status=final_status, zip_url=zip_url )) print(f"[Job {request_id}] ✓ Job completed successfully!") # Log analytics retry_on_network_error(lambda: supabase_client.log_analytics_event( user_id=user_id, event_type="document_generation_completed", entity_id=request_id )) print(f"[Job {request_id}] ✅ Job completed successfully!") except Exception as e: # Update status to failed with error message error_message = f"{type(e).__name__}: {str(e)}" print(f"[Job {request_id}] ❌ Job failed: {error_message}") traceback.print_exc() retry_on_network_error(lambda: supabase_client.update_request_status( request_id=request_id, status="failed", error_message=error_message, zip_url=locals().get('zip_url') )) # Log analytics retry_on_network_error(lambda: supabase_client.log_analytics_event( user_id=user_id, event_type="document_generation_failed", entity_id=request_id )) raise # Re-raise so RQ marks job as failed finally: # Clean up assets directory if it exists if 'assets_temp_dir' in locals() and assets_temp_dir and assets_temp_dir.exists(): try: shutil.rmtree(assets_temp_dir, ignore_errors=True) print(f"[Job {request_id}] ✓ Cleaned up assets directory {assets_temp_dir}") except: pass def process_document_generation_job(request_id: str, request_data: Dict[str, Any]): """ Synchronous wrapper for RQ - calls the async function with asyncio.run(). This is the function that RQ worker calls. It runs the async version using asyncio. """ print(f"{'='*60}") print(f"🎯 Worker picked up job: {request_id}") print(f" User ID: {request_data.get('user_id', 'N/A')}") print(f" Num documents: {request_data.get('prompt_params', {}).get('num_solutions', 'N/A')}") print(f"{'='*60}") return asyncio.run(process_document_generation_job_async(request_id, request_data))