Update app.py

app.py

@@ -6,13 +6,17 @@ import json
 import traceback
 import unicodedata
 import re
+import base64
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
+from typing import Optional
 import fitz  # pymupdf
 from docx import Document
-from docx.shared import Pt, Inches
+from docx.shared import Pt, Inches, RGBColor
 from docx.enum.text import WD_ALIGN_PARAGRAPH
+from docx.oxml.ns import qn
+from PIL import Image
 import requests
 
 app = FastAPI(title="Document Enhancer v5")
@@ -214,14 +218,15 @@ def extract_text_and_layout(file_bytes: bytes, kind: str):
 def enhance_text_with_gemini(text: str, doc_type: str = "auto", user_prompt: str = "") -> str:
     """
     Send prompt to Gemini to produce JSON document layout for ANY document type.
+    Tries multiple API keys if one fails (rate limit, quota exceeded, etc.).
     For very long documents, uses a summarization strategy.
     Returns the raw string response (may include fences), caller should clean it.
     
     doc_type: "auto", "resume", "letter", "report", "article", "essay", etc.
     user_prompt: Custom instructions from user
     """
-    if not GEMINI_API_KEY:
-        return json.dumps({"error": "GEMINI_API_KEY not set"})
+    if not GEMINI_API_KEYS:
+        return json.dumps({"error": "No GEMINI_API_KEY configured"})
 
     # For very long documents, use a more concise output strategy
     is_long_doc = len(text) > 10000
@@ -294,24 +299,61 @@ IMPORTANT:
         }
     }
 
-    try:
-        res = requests.post(GEMINI_URL + f"?key={GEMINI_API_KEY}", json=payload, timeout=240)
-        res.raise_for_status()
-        resp_json = res.json()
-    except Exception as e:
-        return json.dumps({"error": f"gemini request failed: {str(e)}"})
-
-    try:
-        content = resp_json["candidates"][0]["content"]["parts"][0]["text"]
-        
-        # Check if response looks truncated
-        if not content.rstrip().endswith("}") and not content.rstrip().endswith("]"):
-            # Try to repair JSON by closing it
-            content = repair_truncated_json(content)
-        
-        return content
-    except Exception:
-        return json.dumps({"error": resp_json})
+    # Try each API key in sequence until one works
+    last_error = None
+    for idx, api_key in enumerate(GEMINI_API_KEYS):
+        try:
+            print(f"Trying Gemini API key #{idx + 1}/{len(GEMINI_API_KEYS)}...")
+            
+            res = requests.post(
+                GEMINI_URL + f"?key={api_key}", 
+                json=payload, 
+                timeout=240
+            )
+            
+            # Check for rate limit or quota errors
+            if res.status_code == 429:
+                print(f"API key #{idx + 1} rate limited, trying next...")
+                last_error = "Rate limit exceeded"
+                continue
+            
+            if res.status_code == 403:
+                print(f"API key #{idx + 1} quota exceeded or invalid, trying next...")
+                last_error = "Quota exceeded or invalid API key"
+                continue
+            
+            res.raise_for_status()
+            resp_json = res.json()
+            
+            # Extract content
+            content = resp_json["candidates"][0]["content"]["parts"][0]["text"]
+            
+            # Check if response looks truncated
+            if not content.rstrip().endswith("}") and not content.rstrip().endswith("]"):
+                content = repair_truncated_json(content)
+            
+            print(f"✓ API key #{idx + 1} succeeded!")
+            return content
+            
+        except requests.exceptions.Timeout:
+            print(f"API key #{idx + 1} timed out, trying next...")
+            last_error = "Request timed out"
+            continue
+            
+        except requests.exceptions.RequestException as e:
+            print(f"API key #{idx + 1} failed: {str(e)}, trying next...")
+            last_error = str(e)
+            continue
+            
+        except Exception as e:
+            print(f"API key #{idx + 1} error: {str(e)}, trying next...")
+            last_error = str(e)
+            continue
+    
+    # All API keys failed
+    return json.dumps({
+        "error": f"All {len(GEMINI_API_KEYS)} API keys failed. Last error: {last_error}"
+    })
 
 
 # -------------------------
@@ -642,6 +684,93 @@ async def route_full(
         return JSONResponse({"error": str(e)}, status_code=500)
 
 
+@app.post("/add-signature")
+async def route_add_signature(
+    file: UploadFile = File(...),
+    signature: str = Form(...),  # Base64 encoded image
+    position: str = Form("bottom-right"),  # bottom-right, bottom-center, bottom-left, custom
+    signer_name: Optional[str] = Form(None)
+):
+    """
+    Add a signature image to an existing DOCX file.
+    
+    Parameters:
+    - file: DOCX file to sign
+    - signature: Base64 encoded signature image (PNG/JPEG)
+    - position: Where to place signature (bottom-right, bottom-center, bottom-left)
+    - signer_name: Optional name to display under signature
+    """
+    try:
+        # Read the DOCX file
+        docx_bytes = await file.read()
+        doc = Document(io.BytesIO(docx_bytes))
+        
+        # Decode signature image from base64
+        try:
+            # Remove data URI prefix if present (data:image/png;base64,...)
+            if ',' in signature:
+                signature = signature.split(',', 1)[1]
+            
+            signature_bytes = base64.b64decode(signature)
+            signature_image = io.BytesIO(signature_bytes)
+            
+            # Validate it's a valid image
+            img = Image.open(signature_image)
+            signature_image.seek(0)  # Reset stream position
+            
+        except Exception as e:
+            return JSONResponse({"error": f"Invalid signature image: {str(e)}"}, status_code=400)
+        
+        # Get the last section (last page)
+        last_section = doc.sections[-1]
+        
+        # Add signature at the end of document
+        # Add some spacing before signature
+        doc.add_paragraph()
+        
+        # Create signature paragraph
+        sig_paragraph = doc.add_paragraph()
+        
+        # Set alignment based on position
+        if position == "bottom-center":
+            sig_paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER
+        elif position == "bottom-left":
+            sig_paragraph.alignment = WD_ALIGN_PARAGRAPH.LEFT
+        else:  # bottom-right (default)
+            sig_paragraph.alignment = WD_ALIGN_PARAGRAPH.RIGHT
+        
+        # Add signature image (resize to reasonable size)
+        run = sig_paragraph.add_run()
+        run.add_picture(signature_image, width=Inches(1.5))
+        
+        # Add signer name if provided
+        if signer_name and signer_name.strip():
+            name_paragraph = doc.add_paragraph()
+            name_paragraph.alignment = sig_paragraph.alignment
+            name_run = name_paragraph.add_run(signer_name.strip())
+            name_run.font.size = Pt(10)
+            name_run.font.italic = True
+        
+        # Save to bytes
+        output = io.BytesIO()
+        doc.save(output)
+        output.seek(0)
+        
+        safe_filename = sanitize_filename(f"Signed_{file.filename or 'document.docx'}")
+        
+        return StreamingResponse(
+            output,
+            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+            headers={
+                "Content-Disposition": f'attachment; filename="{safe_filename}"'
+            },
+        )
+        
+    except Exception as e:
+        traceback.print_exc()
+        return JSONResponse({"error": f"Failed to add signature: {str(e)}"}, status_code=500)
+
+
 # -------------------------
 # Root
 # -------------------------
@@ -655,6 +784,7 @@ def root():
         "endpoints": {
             "/extract": "Extract plain text from document",
             "/enhance": "Full pipeline: extract + AI enhancement + DOCX (RECOMMENDED)",
+            "/add-signature": "Add signature to existing DOCX file",
             "/preview": "Preview JSON layout without downloading",
             "/design": "Same as /enhance (legacy)",
             "/full": "Same as /enhance (legacy)"
@@ -670,6 +800,7 @@ def root():
 def health():
     return {
         "status": "healthy",
-        "gemini_configured": bool(GEMINI_API_KEY),
-        "version": "5.0 - Universal with Prompt Support"
+        "api_keys_configured": len(GEMINI_API_KEYS),
+        "api_keys_available": GEMINI_API_KEYS[:1] if GEMINI_API_KEYS else [],  # Only show first key for security
+        "version": "5.0 - Universal with Prompt Support + Multi-API-Key Fallback"
     }