Update app.py

app.py

@@ -94,6 +94,42 @@ def is_project_loaded():
     pkl_files = [f for f in os.listdir(extraction_dir) if f.endswith('.pkl')]
     return bool(pkl_files)
 
+# --- Helper Function for Developer Mode ---
+def extract_files_from_response(response):
+    """
+    Parses the LLM response to extract file paths and their corresponding code content.
+
+    Args:
+        response (str): The raw response string from the LLM.
+
+    Returns:
+        dict: A dictionary where keys are file paths and values are the code content of each file.
+    """
+    files = {}
+    current_file = None
+    current_content = []
+
+    for line in response.splitlines():
+        if line.startswith("--- BEGIN FILE:"):
+            if current_file is not None:
+                # Save previous file content
+                files[current_file] = "\n".join(current_content)
+            
+            # Start a new file
+            current_file = line.replace("--- BEGIN FILE:", "").strip()
+            current_content = []
+        elif line.startswith("--- END FILE:"):
+            if current_file is not None:
+                # Save current file content
+                files[current_file] = "\n".join(current_content)
+                current_file = None
+                current_content = []
+        elif current_file is not None:
+            # Append line to current file content
+            current_content.append(line)
+
+    return files
+
 # --- Gradio UI Components ---
 
 # Chat Interface
@@ -122,27 +158,11 @@ def chat_ui(query, history, mode):
         response = "An error occurred during processing. Please check the logs."
 
     if mode == "developer":
-        # Split the response into chunks based on "---"
-        chunks = response.split("---")
-        formatted_response_parts = []
-
-        for chunk in chunks:
-            if chunk.strip().startswith("BEGIN FILE:"):
-                # Extract filepath and code content
-                filepath = chunk.split("BEGIN FILE:")[1].split("\n")[0].strip()
-                code_content = chunk.replace(f"BEGIN FILE: {filepath}\n", "").strip()
-                
-                # Remove "END FILE:" and its associated filepath if present at the end.
-                if code_content.rfind("END FILE:") != -1:
-                    end_file_index = code_content.rfind("END FILE:")
-                    code_content = code_content[:end_file_index].strip()
-                
-                formatted_response_parts.append(f"**{filepath}:**\n`python\n{code_content}\n`")
-            elif chunk.strip():
-                formatted_response_parts.append(chunk.strip())
-
-        # Join the formatted parts with separators
-        formatted_response = "\n---\n".join(formatted_response_parts)
+        extracted_files = extract_files_from_response(response)
+        formatted_response = ""
+
+        for filepath, content in extracted_files.items():
+            formatted_response += f"## {filepath}\n`\n{content}\n`\n"
 
     else:
         # Format the output for non-developer modes