Update app.py

app.py

@@ -73,38 +73,56 @@ def parse_model_response(response_text: str) -> dict:
         if match:
             json_content = match.group(1).strip()
             files = json.loads(json_content)
-            if not isinstance(files, dict) or "files" not in files:
+            if not isinstance(files, dict) or "files" not in files or not isinstance(files["files"], dict):
                 raise ValueError("Invalid JSON structure")
             socketio.emit('progress', {'percent': 90, 'message': 'JSON parsed'})
             return files
         else:
             app.logger.warning("No JSON found within START/STOP tokens, attempting custom parsing")
-            # Fallback to the existing regex parsing
             pattern = r"### (.+?)\n```(?:\w+)?\n(.*?)\n```"
             matches = re.findall(pattern, response_text, re.DOTALL)
             if not matches:
-                raise ValueError(f"Could not parse response into files. Response start: {response_text[:200]}")  # Log response start
+                raise ValueError(f"Could not parse response into files. Response start: {response_text[:200]}")
             files = {
                 "files": {
                     filename.strip(): {"content": content.strip()}
                     for filename, content in matches
                 }
             }
+            # Extract body content from index.html if present
+            index_html_content = files["files"].get("index.html", {}).get("content", "")
+            body_match = re.search(r"<body.*?>(.*?)</body>", index_html_content, re.DOTALL | re.IGNORECASE)
+            if body_match:
+                files["files"]["index.html"]["content"] = body_match.group(1).strip()
+                app.logger.info("Extracted body content from index.html")
+            else:
+                app.logger.warning("Could not extract body content from index.html, using full content.")
+
+
             socketio.emit('progress', {'percent': 90, 'message': 'Parsed with fallback method'})
             return files
 
     except json.JSONDecodeError:
-        app.logger.warning(f"Response is not valid JSON, attempting custom parsing. Response start: {response_text[:200]}")  # Log response start
+        app.logger.warning(f"Response is not valid JSON, attempting custom parsing. Response start: {response_text[:200]}")
         pattern = r"### (.+?)\n```(?:\w+)?\n(.*?)\n```"
         matches = re.findall(pattern, response_text, re.DOTALL)
         if not matches:
-            raise ValueError(f"Could not parse response into files. Response start: {response_text[:200]}")  # Log response start
+            raise ValueError(f"Could not parse response into files. Response start: {response_text[:200]}")
         files = {
             "files": {
                 filename.strip(): {"content": content.strip()}
-                for filename, content in matches
+                    for filename, content in matches
             }
         }
+        # Extract body content from index.html if present
+        index_html_content = files["files"].get("index.html", {}).get("content", "")
+        body_match = re.search(r"<body.*?>(.*?)</body>", index_html_content, re.DOTALL | re.IGNORECASE)
+        if body_match:
+            files["files"]["index.html"]["content"] = body_match.group(1).strip()
+            app.logger.info("Extracted body content from index.html")
+        else:
+            app.logger.warning("Could not extract body content from index.html, using full content.")
+
         socketio.emit('progress', {'percent': 90, 'message': 'Parsed with fallback method'})
         return files
     except ValueError as e:  # Catch ValueErrors too