Update app.py

app.py

@@ -624,123 +624,207 @@ def extract_json_from_llm_response(raw_response: str) -> dict:
         logger.error("Sanitized JSON still invalid:\n%s", json_string)
         raise
 
-# def clean_base64_for_model(raw_b64):
+# def reduce_image_size_to_limit(clean_b64_str, max_kb=4000):
 #     """
-#     Normalize input into a valid data:image/png;base64,<payload> string.
-
-#     Accepts:
-#       - a list of base64 strings → picks the first element
-#       - a PIL Image instance       → encodes to PNG/base64
-#       - a raw base64 string        → strips whitespace and data URI prefix
+#     Reduce an image's size to be as close as possible to max_kb without exceeding it.
+#     Returns the final base64 string and its size in KB.
 #     """
+#     import re, base64
+#     from io import BytesIO
+#     from PIL import Image
+
+#     # Remove the data URI prefix
+#     base64_data = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", clean_b64_str)
+#     image_data = base64.b64decode(base64_data)
+
+#     # Load into PIL
+#     img = Image.open(BytesIO(image_data))
+
+#     low, high = 20, 95  # reasonable JPEG quality range
+#     best_b64 = None
+#     best_size_kb = 0
+
+#     while low <= high:
+#         mid = (low + high) // 2
+#         buffer = BytesIO()
+#         img.save(buffer, format="JPEG", quality=mid)
+#         size_kb = len(buffer.getvalue()) / 1024
+
+#         if size_kb <= max_kb:
+#             # This quality is valid, try higher
+#             best_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+#             best_size_kb = size_kb
+#             low = mid + 1
+#         else:
+#             # Too big, try lower
+#             high = mid - 1
+
+#     return f"data:image/jpeg;base64,{best_b64}"
+
+# #clean the base64 model here
+# def clean_base64_for_model(raw_b64):
+#     import io, base64, re
+#     from PIL import Image
+
 #     if not raw_b64:
-#         return ""
+#         return "", ""
 
-#     # 1. If it’s a list, take its first element
 #     if isinstance(raw_b64, list):
 #         raw_b64 = raw_b64[0] if raw_b64 else ""
 #         if not raw_b64:
-#             return ""
+#             return "", ""
 
-#     # 2. If it’s a PIL Image, convert to base64
 #     if isinstance(raw_b64, Image.Image):
 #         buf = io.BytesIO()
 #         raw_b64.save(buf, format="PNG")
 #         raw_b64 = base64.b64encode(buf.getvalue()).decode()
 
-#     # 3. At this point it must be a string
 #     if not isinstance(raw_b64, str):
 #         raise TypeError(f"Expected base64 string or PIL Image, got {type(raw_b64)}")
 
-#     # 4. Strip any existing data URI prefix, whitespace, or newlines
+#     # Remove data URI prefix if present
 #     clean_b64 = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", raw_b64)
 #     clean_b64 = clean_b64.replace("\n", "").replace("\r", "").strip()
 
-#     # 5. Validate it’s proper base64
-#     try:
-#         base64.b64decode(clean_b64)
-#     except Exception as e:
-#         logger.error(f"Invalid Base64 passed to model: {e}")
-#         raise
-
-#     # 6. Return with the correct data URI prefix
-#     return f"data:image/png;base64,{clean_b64}"
-
-def reduce_image_size_to_limit(clean_b64_str, max_kb=4000):
+#     # Log original size
+#     original_size = len(clean_b64.encode("utf-8"))
+#     print(f"Original Base64 size (bytes): {original_size}")
+#     if original_size > 4000000:
+#         # Reduce size to under 4 MB
+#         reduced_b64 = reduce_image_size_to_limit(clean_b64, max_kb=4000)
+#         clean_b64_2 = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", reduced_b64)
+#         clean_b64_2 = clean_b64_2.replace("\n", "").replace("\r", "").strip()
+#         reduced_size = len(clean_b64_2.encode("utf-8"))
+#         print(f"Reduced Base64 size (bytes): {reduced_size}")
+#         # Return both prefixed and clean reduced versions
+#         return f"data:image/jpeg;base64,{reduced_b64}"
+#     return f"data:image/jpeg;base64,{clean_b64}"
+
+def reduce_image_size_to_limit(clean_b64_str: str, max_kb: int = 4000) -> str:
     """
-    Reduce an image's size to be as close as possible to max_kb without exceeding it.
-    Returns the final base64 string and its size in KB.
+    Input: clean_b64_str = BASE64 STRING (no data: prefix)
+    Output: BASE64 STRING (no data: prefix), sized as close as possible to max_kb KB.
+    Guarantees: returns a valid base64 string (never None). May still be larger than max_kb
+    if saving at lowest quality cannot get under the limit.
     """
-    import re, base64
-    from io import BytesIO
-    from PIL import Image
+    # sanitize
+    clean = re.sub(r"\s+", "", clean_b64_str).strip()
+    # fix padding
+    missing = len(clean) % 4
+    if missing:
+        clean += "=" * (4 - missing)
 
-    # Remove the data URI prefix
-    base64_data = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", clean_b64_str)
-    image_data = base64.b64decode(base64_data)
-
-    # Load into PIL
-    img = Image.open(BytesIO(image_data))
-
-    low, high = 20, 95  # reasonable JPEG quality range
-    best_b64 = None
-    best_size_kb = 0
+    try:
+        image_data = base64.b64decode(clean)
+    except Exception as e:
+        raise ValueError("Invalid base64 input to reduce_image_size_to_limit") from e
 
+    try:
+        img = Image.open(io.BytesIO(image_data))
+        img.load()
+    except Exception as e:
+        raise ValueError("Could not open image from base64") from e
+
+    # convert alpha -> RGB because JPEG doesn't support alpha
+    if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info):
+        background = Image.new("RGB", img.size, (255, 255, 255))
+        background.paste(img, mask=img.split()[-1] if img.mode != "RGB" else None)
+        img = background
+    elif img.mode != "RGB":
+        img = img.convert("RGB")
+
+    low, high = 20, 95
+    best_bytes = None
+    # binary search for best quality
     while low <= high:
         mid = (low + high) // 2
-        buffer = BytesIO()
-        img.save(buffer, format="JPEG", quality=mid)
-        size_kb = len(buffer.getvalue()) / 1024
-
+        buf = io.BytesIO()
+        try:
+            img.save(buf, format="JPEG", quality=mid, optimize=True)
+        except OSError:
+            # some PIL builds/channels may throw on optimize=True; fallback without optimize
+            buf = io.BytesIO()
+            img.save(buf, format="JPEG", quality=mid)
+        size_kb = len(buf.getvalue()) / 1024.0
         if size_kb <= max_kb:
-            # This quality is valid, try higher
-            best_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
-            best_size_kb = size_kb
+            best_bytes = buf.getvalue()
             low = mid + 1
         else:
-            # Too big, try lower
             high = mid - 1
 
-    return f"data:image/jpeg;base64,{best_b64}"
+    # if never found a quality <= max_kb, use the smallest we created (quality = 20)
+    if best_bytes is None:
+        buf = io.BytesIO()
+        try:
+            img.save(buf, format="JPEG", quality=20, optimize=True)
+        except OSError:
+            buf = io.BytesIO()
+            img.save(buf, format="JPEG", quality=20)
+        best_bytes = buf.getvalue()
 
-#clean the base64 model here
-def clean_base64_for_model(raw_b64):
-    import io, base64, re
-    from PIL import Image
+    return base64.b64encode(best_bytes).decode("utf-8")
 
+
+def clean_base64_for_model(raw_b64, max_bytes_threshold=4000000) -> str:
+    """
+    Accepts: raw_b64 can be:
+      - a data URI 'data:image/png;base64,...'
+      - a plain base64 string
+      - a PIL Image
+      - a list containing the above (take first)
+    Returns: a data URI string 'data:<mime>;base64,<base64>' guaranteed to be syntactically valid.
+    """
+    # normalize input
     if not raw_b64:
-        return "", ""
+        return ""
 
     if isinstance(raw_b64, list):
         raw_b64 = raw_b64[0] if raw_b64 else ""
         if not raw_b64:
-            return "", ""
+            return ""
 
     if isinstance(raw_b64, Image.Image):
         buf = io.BytesIO()
-        raw_b64.save(buf, format="PNG")
-        raw_b64 = base64.b64encode(buf.getvalue()).decode()
+        # convert to RGB and save as JPEG to keep consistent
+        img = raw_b64.convert("RGB")
+        img.save(buf, format="JPEG")
+        clean_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        mime = "image/jpeg"
+        return f"data:{mime};base64,{clean_b64}"
 
     if not isinstance(raw_b64, str):
         raise TypeError(f"Expected base64 string or PIL Image, got {type(raw_b64)}")
 
-    # Remove data URI prefix if present
-    clean_b64 = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", raw_b64)
-    clean_b64 = clean_b64.replace("\n", "").replace("\r", "").strip()
-
-    # Log original size
-    original_size = len(clean_b64.encode("utf-8"))
-    print(f"Original Base64 size (bytes): {original_size}")
-    if original_size > 4000000:
-        # Reduce size to under 4 MB
-        reduced_b64 = reduce_image_size_to_limit(clean_b64, max_kb=4000)
-        clean_b64_2 = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", reduced_b64)
-        clean_b64_2 = clean_b64_2.replace("\n", "").replace("\r", "").strip()
-        reduced_size = len(clean_b64_2.encode("utf-8"))
-        print(f"Reduced Base64 size (bytes): {reduced_size}")
-        # Return both prefixed and clean reduced versions
-        return f"data:image/jpeg;base64,{reduced_b64}"
-    return f"data:image/jpeg;base64,{clean_b64}"
+    # detect mime if present; otherwise default to png
+    m = re.match(r"^data:(image\/[a-zA-Z0-9.+-]+);base64,(.+)$", raw_b64, flags=re.DOTALL)
+    if m:
+        mime = m.group(1)
+        clean_b64 = m.group(2)
+    else:
+        # no prefix; assume png by default (you can change to jpeg if you prefer)
+        mime = "image/png"
+        clean_b64 = raw_b64
+
+    # sanitize base64 string
+    clean_b64 = re.sub(r"\s+", "", clean_b64).strip()
+    missing = len(clean_b64) % 4
+    if missing:
+        clean_b64 += "=" * (4 - missing)
+
+    original_size_bytes = len(clean_b64.encode("utf-8"))
+    # debug print
+    print(f"Original base64 size (bytes): {original_size_bytes}, mime: {mime}")
+
+    if original_size_bytes > max_bytes_threshold:
+        # reduce and return JPEG prefixed data URI (JPEG tends to compress better for photos)
+        reduced_clean = reduce_image_size_to_limit(clean_b64, max_kb=4000)
+        # reduced_clean is plain base64 (no prefix)
+        print(f"Reduced base64 size (bytes): {original_size_bytes}, mime: {mime}")
+        return f"data:image/jpeg;base64,{reduced_clean}"
+
+    # otherwise return original with its mime prefix (ensure prefix exists)
+    return f"data:{mime};base64,{clean_b64}"
+
 
 SCRATCH_OPCODES = [
     'motion_movesteps', 'motion_turnright', 'motion_turnleft', 'motion_goto',