Update app.py

app.py

@@ -1484,9 +1484,25 @@ async def enhance_document(
     prompt: str = Form(...),
     output_format: str = Form("docx"),
 ):
-    logger.info("Starting enhancement for file: %s", file.filename)
+    """
+    Enhanced document processing workflow:
+    1. OCR Extract - Extract text and layout from uploaded file
+    2. Cleanup - Remove headers/footers and artifacts
+    3. Gemini Enhancement - Enhance text based on user prompt
+    4. Gemini Layout - Generate best layout for enhanced text
+    5. DOCX Build - Create final formatted document
+    """
+    logger.info("=" * 70)
+    logger.info("📄 STARTING DOCUMENT ENHANCEMENT WORKFLOW")
+    logger.info("=" * 70)
+    logger.info(f"File: {file.filename} | Prompt: {prompt[:50]}...")
+
     start_time = asyncio.get_event_loop().time()
 
+    # ============ STEP 1: FILE VALIDATION & EXTRACTION ============
+    logger.info("\n[STEP 1/5] 📂 FILE EXTRACTION")
+    logger.info("-" * 70)
+
     raw = _read_bytes(file)
     if not raw:
         raise HTTPException(400, detail="Empty file")
@@ -1494,113 +1510,131 @@ async def enhance_document(
     file_kind = _detect_file_kind(file.filename, file.content_type)
     suffix = os.path.splitext(file.filename or "upload.bin")[1] or ".bin"
     temp_path = _save_to_temp(raw, suffix=suffix)
-    logger.info(f"File saved: {file_kind}, size: {len(raw)} bytes")
+    logger.info(f"✓ File saved: {file_kind} | Size: {len(raw):,} bytes")
 
     try:
-        # Step 1: Extract text (fast)
-        plain_text, _ = extract_text_and_layout(temp_path, file_kind)
+        # ============ STEP 2: OCR & TEXT EXTRACTION ============
+        logger.info("\n[STEP 2/5] 🔍 OCR TEXT EXTRACTION")
+        logger.info("-" * 70)
+
+        plain_text, layout_info = extract_text_and_layout(temp_path, file_kind)
         if not plain_text.strip():
             raise HTTPException(422, detail="Could not extract text from the document")
 
-        logger.info(f"Extracted text: {len(plain_text)} chars")
+        logger.info(f"✓ Text extracted: {len(plain_text):,} characters")
+        if layout_info:
+            logger.info(f"✓ Layout info extracted: {len(layout_info):,} bytes")
+
+        # ============ STEP 3: CLEANUP & PREPROCESSING ============
+        logger.info("\n[STEP 3/5] 🧹 CLEANUP & PREPROCESSING")
+        logger.info("-" * 70)
 
-        # Remove headers/footers using fast pattern detection
-        logger.info("=" * 60)
-        logger.info("🧹 HEADER/FOOTER REMOVAL STARTING...")
-        logger.info("=" * 60)
         original_len = len(plain_text)
         plain_text = _remove_headers_footers_fast(plain_text)
         removed_pct = (
             100 * (1 - len(plain_text) / original_len) if original_len > 0 else 0
         )
-        logger.info("=" * 60)
         logger.info(
-            f"✓ CLEANING COMPLETE: {len(plain_text)} chars (removed {removed_pct:.1f}%)"
+            f"✓ Cleanup complete: {len(plain_text):,} characters (removed {removed_pct:.1f}%)"
         )
-        logger.info("=" * 60)
 
-        # For very large documents, use fast mode (skip AI enhancement, use smart layout)
-        if len(plain_text) > 25000:
-            logger.info(
-                f"Very large document ({len(plain_text)} chars) - using FAST MODE"
-            )
+        # ============ STEP 4: GEMINI TEXT ENHANCEMENT ============
+        logger.info("\n[STEP 4/5] ✨ GEMINI TEXT ENHANCEMENT")
+        logger.info("-" * 70)
+        logger.info(f"User prompt: {prompt}")
+
+        # Validate user prompt
+        if not prompt or prompt.strip().lower() in ("", "enhance", "improve"):
+            logger.info("⚠ No specific prompt provided, using default enhancement")
+            enhancement_prompt = "Improve this document by enhancing clarity, fixing grammar, improving flow, and using professional language while preserving all content."
+        else:
+            enhancement_prompt = prompt
+            logger.info("✓ Using custom enhancement prompt")
+
+        # Process text enhancement (with chunking for large docs)
+        if len(plain_text) > 10000:
             logger.info(
-                "Skipping AI enhancement to avoid timeout, creating smart layout..."
+                f"⚠ Large document ({len(plain_text):,} chars) - using chunked enhancement"
             )
+            chunks = _chunk_text(plain_text, max_chars=10000)
+            logger.info(f"  Split into {len(chunks)} chunks")
+
+            enhanced_chunks = []
+            for i, chunk in enumerate(chunks):
+                logger.info(f"  Processing chunk {i + 1}/{len(chunks)}...")
+                try:
+                    enhanced = await gemini_edit_text(
+                        prompt=enhancement_prompt, text=chunk, timeout=120
+                    )
+                    enhanced_chunks.append(enhanced)
+                    logger.info(
+                        f"    ✓ Chunk {i + 1} enhanced: {len(enhanced):,} chars"
+                    )
+                except Exception as e:
+                    logger.warning(
+                        f"    ✗ Chunk {i + 1} failed: {str(e)[:100]} - using original"
+                    )
+                    enhanced_chunks.append(chunk)
 
-            # Create professional layout without AI
-            design = _create_smart_layout(plain_text)
-            content_text = plain_text
+            enhanced_text = "\n\n".join(enhanced_chunks)
         else:
-            # Step 2: Edit text with timeout and chunking for medium docs
-            logger.info("Starting text enhancement...")
+            logger.info(f"Processing {len(plain_text):,} characters...")
+            enhanced_text = await gemini_edit_text(
+                prompt=enhancement_prompt, text=plain_text, timeout=120
+            )
 
-            # Validate and enhance the user prompt
-            if not prompt or prompt.strip().lower() in ("", "enhance", "improve"):
-                logger.info(
-                    "No specific enhancement instruction provided, using default enhancement"
-                )
-                user_enhancement_prompt = "Improve this document by enhancing clarity, fixing grammar, improving flow, and using professional language while preserving all content."
-            else:
-                user_enhancement_prompt = prompt
-                logger.info(
-                    f"Using user enhancement instruction: {user_enhancement_prompt[:100]}..."
-                )
+        # Verify enhancement
+        if len(enhanced_text.strip()) == 0:
+            logger.warning("⚠ Enhancement resulted in empty text - using original")
+            enhanced_text = plain_text
 
-            if len(plain_text) > 10000:
-                logger.info("Medium-large document detected, using chunking strategy")
-                chunks = _chunk_text(plain_text, max_chars=10000)
-                logger.info(f"Split into {len(chunks)} chunks")
-                enhanced_chunks = []
-                for i, chunk in enumerate(chunks):
-                    logger.info(f"Processing chunk {i + 1}/{len(chunks)}")
-                    try:
-                        enhanced = await gemini_edit_text(
-                            prompt=user_enhancement_prompt, text=chunk, timeout=120
-                        )
-                        enhanced_chunks.append(enhanced)
-                    except Exception as e:
-                        logger.warning(
-                            f"Chunk {i + 1} failed: {str(e)}, using original"
-                        )
-                        enhanced_chunks.append(chunk)
-                content_text = "\n\n".join(enhanced_chunks)
-            else:
-                content_text = await gemini_edit_text(
-                    prompt=user_enhancement_prompt, text=plain_text, timeout=120
-                )
+        enhancement_ratio = (
+            len(enhanced_text) / len(plain_text) if len(plain_text) > 0 else 1.0
+        )
+        logger.info(
+            f"✓ Text enhanced: {len(enhanced_text):,} characters (ratio: {enhancement_ratio:.2f}x)"
+        )
 
-            # Verify enhancement actually happened
-            if len(content_text.strip()) == 0:
-                logger.warning("Enhancement resulted in empty text, using original")
-                content_text = plain_text
+        # ============ STEP 5: GEMINI LAYOUT GENERATION ============
+        logger.info("\n[STEP 5/5] 🎨 GEMINI LAYOUT GENERATION")
+        logger.info("-" * 70)
 
-            logger.info(
-                f"Enhanced text: {len(content_text)} chars (original: {len(plain_text)} chars)"
-            )
+        # Use reasonable amount for design (up to 20K chars for better layouts)
+        design_input_text = (
+            enhanced_text[:20000] + "\n\n[Content continues...]"
+            if len(enhanced_text) > 20000
+            else enhanced_text
+        )
+        logger.info(f"Generating layout for {len(design_input_text):,} characters...")
 
-            # Step 3: Generate design with timeout
-            logger.info("Starting design generation...")
-            instruction = _house_style_instruction_v2()
-            # Use reasonable amount for design (up to 15K chars)
-            design_text = (
-                content_text[:15000] + "\n\n[Content continues...]"
-                if len(content_text) > 15000
-                else content_text
-            )
-            logger.info(f"Design generation using {len(design_text)} chars")
+        try:
             design = await gemini_generate_design(
-                user_instruction=instruction, source_text=design_text
+                user_instruction="", source_text=design_input_text
+            )
+            logger.info(f"✓ Layout generated: {len(design.get('document', []))} blocks")
+        except Exception as e:
+            logger.warning(
+                f"⚠ Layout generation failed: {str(e)[:100]} - using smart layout"
+            )
+            design = _create_smart_layout(enhanced_text)
+            logger.info(
+                f"✓ Smart layout applied: {len(design.get('document', []))} blocks"
             )
 
-        # Step 4: Build DOCX (fast)
-        logger.info("Building DOCX...")
+        # ============ STEP 6: DOCX BUILDING ============
+        logger.info("\n[STEP 6/5] 📝 BUILDING DOCX OUTPUT")
+        logger.info("-" * 70)
+
         out_bytes = build_docx_from_design(design)
+        logger.info(f"✓ DOCX built: {len(out_bytes):,} bytes")
 
+        # ============ COMPLETION ============
         elapsed = asyncio.get_event_loop().time() - start_time
-        logger.info(
-            f"✓ SUCCESS - Total time: {elapsed:.2f}s, DOCX size: {len(out_bytes)} bytes"
-        )
+        logger.info("\n" + "=" * 70)
+        logger.info(f"✅ SUCCESS - Total time: {elapsed:.2f}s")
+        logger.info(f"   Input: {len(raw):,} bytes → Output: {len(out_bytes):,} bytes")
+        logger.info(f"   Text: {len(plain_text):,} → {len(enhanced_text):,} chars")
+        logger.info("=" * 70 + "\n")
 
         out_name_base = (
             os.path.splitext(os.path.basename(file.filename or "document"))[0]
@@ -1613,18 +1647,22 @@ async def enhance_document(
             media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
             headers={"Content-Disposition": f'attachment; filename="{out_name}"'},
         )
+
     except asyncio.TimeoutError:
-        logger.error("✗ Request timed out after waiting too long")
+        logger.error("❌ TIMEOUT - Request took too long")
         raise HTTPException(
             504,
             detail="Request timed out. Document may be too large or complex. Try splitting it into smaller parts.",
         )
+    except HTTPException:
+        raise
     except Exception as e:
-        logger.error(f"✗ Enhancement failed: {str(e)}")
+        logger.error(f"❌ FAILED - {str(e)}")
         raise HTTPException(500, detail=f"Processing failed: {str(e)}")
     finally:
         try:
             os.remove(temp_path)
+            logger.info("✓ Temporary file cleaned up")
         except Exception:
             pass