Update app.py

app.py

@@ -435,21 +435,35 @@ 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: str) -> str:
+def clean_base64_for_model(raw_b64):
     if not raw_b64:
         return ""
-    # Remove prefix if exists
+    
+    # If it's a list, take the first image (or join if multiple are expected)
+    if isinstance(raw_b64, list):
+        if len(raw_b64) == 0:
+            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
+    if not isinstance(raw_b64, str):
+        raise TypeError(f"Expected base64 string, got {type(raw_b64)}")
+    
+    # Remove prefix
     clean_b64 = re.sub(r"^data:image\/[a-zA-Z]+;base64,", "", raw_b64)
-    # Remove spaces/newlines
     clean_b64 = clean_b64.replace("\n", "").replace("\r", "").strip()
+    
     # Validate
     try:
         base64.b64decode(clean_b64)
     except Exception as e:
         logger.error(f"Invalid Base64 passed to model: {e}")
         raise
+    
     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 ---")