Update main.py

main.py

@@ -240,38 +240,43 @@ def generate_master_blueprint_task(subject, flattened_data, uid, epiphany_id):
         return None
 
 #Prepare hi fidelity images to proper scale
-def prepare_vision_bytes(image_bytes, max_size_mb=9.5):
+def prepare_vision_image(image_bytes):
     """
-    Surgically compresses/resizes image bytes to stay under the 10MB API limit.
-    Target: 9.5MB to allow a safe margin for the rest of the request payload.
+    Resizes and optimizes the image to ensure it is under the 10MB Gemini limit.
+    Returns a PIL Image object ready for the SDK.
     """
-    size_mb = len(image_bytes) / (1024 * 1024)
-    if size_mb <= max_size_mb:
-        return image_bytes
-
-    logger.info(f"Vision Scaler: Image is {size_mb:.2f}MB. Optimizing for API...")
     try:
         img = Image.open(io.BytesIO(image_bytes))
         
-        # Target resolution: 2048px on the long side is optimal for Gemini detection
-        max_dim = 2048
+        # 1. Convert to RGB to strip Alpha channels (saves 25% space)
+        if img.mode != 'RGB':
+            img = img.convert('RGB')
+
+        # 2. If the file is huge, reduce the physical resolution.
+        # 2000px is more than enough for Athena to see components.
+        max_dimension = 2000
         w, h = img.size
-        if max(w, h) > max_dim:
-            scale = max_dim / max(w, h)
-            img = img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
-
-        output = io.BytesIO()
-        # Convert to RGB (removes Alpha channel if present to save space)
-        # Save as optimized JPEG at 85% quality
-        img.convert('RGB').save(output, format="JPEG", quality=85, optimize=True)
-        new_bytes = output.getvalue()
+        if max(w, h) > max_dimension:
+            scale = max_dimension / max(w, h)
+            new_size = (int(w * scale), int(h * scale))
+            img = img.resize(new_size, Image.Resampling.LANCZOS)
+            logger.info(f"Theia Vision: Resized from {w}x{h} to {new_size[0]}x{new_size[1]}")
+
+        # 3. Final Quality Check: If the image is still dense, 
+        # we do an in-memory compression cycle to strip metadata.
+        buffer = io.BytesIO()
+        img.save(buffer, format="JPEG", quality=85, optimize=True)
+        final_img = Image.open(buffer)
+        
+        final_size = len(buffer.getvalue())
+        logger.info(f"Theia Vision: Optimized payload size: {final_size / 1024 / 1024:.2f}MB")
+        
+        return final_img, final_size
         
-        new_size_mb = len(new_bytes)/(1024*1024)
-        logger.info(f"Vision Scaler: Optimization complete. New size: {new_size_mb:.2f}MB")
-        return new_bytes
     except Exception as e:
-        logger.error(f"Vision Scaler Failure: {e}")
-        return image_bytes # Fallback to original and hope for the best
+        logger.error(f"Theia Vision Scaler Failure: {e}")
+        # Fallback to original
+        return Image.open(io.BytesIO(image_bytes))
 # -----------------------------------------------------------------------------
 # 4. PRIMARY ENDPOINTS: GENERATE & THEIA SWEEP
 # -----------------------------------------------------------------------------
@@ -305,8 +310,11 @@ def generate_epiphany():
         return jsonify({'error': 'Visual image is required.'}), 400
     
     image_file = request.files['image']
-    image_bytes = image_file.read()
+    raw_bytes = image_file.read()
+     # THE FIX:
+    image_bytes = prepare_vision_bytes(raw_bytes)
     pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
+    
 
     try:
         # Step 1: Accurate Identification
@@ -410,7 +418,7 @@ def generate_epiphany():
 
 @app.route('/api/epiphany/theia', methods=['POST'])
 def theia_sweep():
-    """Standalone Theia Mode: Bounding Box Annotations with Vision Scaling."""
+    """Standalone Theia Mode: Bounding Box Annotations with Automatic Image Scaling."""
     logger.info(">>> THEIA SWEEP INITIATED")
     uid = verify_token(request.headers.get('Authorization'))
     if not uid: return jsonify({'error': 'Unauthorized'}), 401
@@ -418,27 +426,32 @@ def theia_sweep():
     epiphany_id = request.form.get('epiphanyId')
     if not epiphany_id: return jsonify({'error': 'epiphanyId is required.'}), 400
 
-    # 1. Check if annotations already exist to prevent double-charging
+    # 1. Cache Check
     epiphany_ref = db_ref.child(f'epiphanies/{epiphany_id}')
     existing_data = epiphany_ref.get() or {}
     if 'annotations' in existing_data:
-        return jsonify({"annotations": existing_data['annotations'], "status": "already_stored"}), 200
+        logger.info(f"Theia: Returning cached annotations for {epiphany_id}")
+        return jsonify({"annotations": existing_data['annotations'], "status": "cached"}), 200
 
-    # 2. Check Sparks
+    # 2. Credit Check
     user_ref = db_ref.child(f'users/{uid}')
     user_data = user_ref.get() or {}
     if user_data.get('credits', 0) < 4:
         return jsonify({'error': 'Need 4 Sparks for a Theia Sweep.'}), 402
 
-    # --- IMAGE PRE-PROCESSING ---
+    # 3. Image Handling with Scaler
     if 'image' not in request.files:
         return jsonify({'error': 'image file is required.'}), 400
         
     image_file = request.files['image']
     raw_bytes = image_file.read()
     
-    # Surgical Fix: Stay under 10MB limit
-    image_bytes = prepare_vision_bytes(raw_bytes)
+    # --- THE CRITICAL FIX ---
+    # We process the raw bytes into a scaled/optimized PIL Image
+    
+    pil_image, scaled_size = prepare_vision_image(raw_bytes)
+    
+    logger.info(f"Theia Post-Scale Size: {scaled_size / (1024 * 1024):.2f} MB")
     
     subject = existing_data.get('subject', 'Complex System')
 
@@ -453,7 +466,7 @@ def theia_sweep():
     """
 
     try:
-        pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
+        # Pass the optimized PIL image to the SDK
         res = client.models.generate_content(
             model=ATHENA_FLASH,
             contents=[sweep_prompt, pil_image],
@@ -464,19 +477,21 @@ def theia_sweep():
         )
         
         raw_json = res.text.strip()
-        # Use helper to strip fences if present
-        if "```" in raw_json:
-            raw_json = _strip_json_fences(raw_json)
+        if "```json" in raw_json:
+            raw_json = re.search(r'```json\n(.*?)\n```', raw_json, re.DOTALL).group(1)
+        elif "```" in raw_json:
+            raw_json = raw_json.replace("```", "").strip()
             
         annotations = json.loads(raw_json)
         
-        # 3. STORE IN DB & CHARGE
+        # 4. Persistence & Deduction
         epiphany_ref.update({"annotations": annotations})
         user_ref.update({'credits': user_data.get('credits', 0) - 4})
         
+        logger.info(f"THEIA SUCCESS: {len(annotations)} annotations found.")
         return jsonify({"annotations": annotations}), 200
     except Exception as e:
-        logger.error(f"Theia Sweep Error: {e}")
+        logger.error(f"Theia Sweep Execution Error: {e}")
         return jsonify({'error': str(e)}), 500
 
 @app.route('/api/epiphany/deep-dive', methods=['POST'])