Update app.py

app.py

@@ -93,6 +93,7 @@ def is_project_loaded():
     extraction_dir = "extraction"
     pkl_files = [f for f in os.listdir(extraction_dir) if f.endswith('.pkl')]
     return bool(pkl_files)
+
 # --- Gradio UI Components ---
 
 # Chat Interface
@@ -101,7 +102,7 @@ def chat_ui(query, history, mode):
     api_key = load_api_key()
     if not api_key:
         return "Error: OpenAI API key not set. Please set the API key in the Settings tab.", []
-    
+
     if not is_project_loaded():
         return "Error: No project loaded. Please upload and process a ZIP file first.", []
 
@@ -121,59 +122,35 @@ def chat_ui(query, history, mode):
         response = "An error occurred during processing. Please check the logs."
 
     if mode == "developer":
-        extracted_files = extract_files_from_response(response)
-
-        # Format the output for developer mode
-        developer_response = ""
-        for filepath, content in extracted_files.items():
-            developer_response += f"**{filepath}:**\n`python\n{content}\n`\n\n" 
-
-        history.append((query, developer_response))
-        # Return history and an empty string for the text output (as it's handled by the chatbot)
-        return history, history  
+        # 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)
 
     else:
         # Format the output for non-developer modes
-        formatted_response = response.replace('\n', '  \n')  # Use two spaces for markdown line breaks
-        history.append((query, formatted_response))
-        # Return history and an empty string for the text output (as it's handled by the chatbot)
-        return history, history  
-
-
-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
+        formatted_response = response.replace('\n', '  \n')
+
+    history.append((query, formatted_response))
+
+    return history, history
 
 # ZIP Processing Interface
 zip_iface = gr.Interface(