Update app.py

app.py

@@ -435,35 +435,55 @@ def extract_json_from_llm_response(raw_response: str) -> dict:
         logger.error("Sanitized JSON still invalid:\n%s", json_string)
         raise
 
+import re
+import io
+import base64
+from PIL import Image
+
 def clean_base64_for_model(raw_b64):
+    """
+    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
+    """
     if not raw_b64:
         return ""
-    
-    # If it's a list, take the first image (or join if multiple are expected)
+
+    # 1. If it’s a list, take its first element
     if isinstance(raw_b64, list):
-        if len(raw_b64) == 0:
+        raw_b64 = raw_b64[0] if raw_b64 else ""
+        if not raw_b64:
             return ""
-        raw_b64 = raw_b64[0]  # take first element
-        # Or: raw_b64 = raw_b64[-1]  # take last element if that’s your logic
-    
-    # Ensure string
+
+    # 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, got {type(raw_b64)}")
-    
-    # Remove prefix
+        raise TypeError(f"Expected base64 string or PIL Image, got {type(raw_b64)}")
+
+    # 4. Strip any existing data URI prefix, whitespace, or newlines
     clean_b64 = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", raw_b64)
     clean_b64 = clean_b64.replace("\n", "").replace("\r", "").strip()
-    
-    # Validate
+
+    # 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}"
 
 
+
 # Node 1: Logic updating if any issue here
 def pseudo_generator_node(state: GameState):
     logger.info("--- Running plan_logic_aligner_node ---")