Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import gradio as gr
 from crewai import Agent, Task, Crew, Process, LLM
-from crewai_tools import FileWriterTool
+from crewai_tools import FileWriterTool, FileReadTool
 from dotenv import load_dotenv
 import traceback
 import shutil
@@ -479,7 +479,7 @@ Target Audience: {target_audience}
 # --- SYSTEM SETUP ---
 try:
     llm = LLM(
-        model="gemini/gemini-2.5-pro",
+        model="gemini/gemini-pro",
         temperature=0.7,
         api_key=os.environ.get("GOOGLE_API_KEY")
     )
@@ -489,6 +489,8 @@ except Exception as e:
 
 os.makedirs("outputs", exist_ok=True)
 file_write_tool = FileWriterTool()
+# Add FileReadTool for modification tasks
+file_read_tool = FileReadTool()
 
 # --- ENHANCED AGENTS ---
 designer = Agent(
@@ -572,7 +574,7 @@ developer = Agent(
     ),
     verbose=True,
     llm=llm,
-    tools=[file_write_tool],
+    tools=[file_write_tool, file_read_tool], # Give agent both read and write tools
     allow_delegation=False
 )
 
@@ -699,7 +701,7 @@ def create_development_task(website_type, page_type):
             "- Implement smooth scrolling\n"
             "- Add hover effects\n"
             "- Include mobile menu\n\n"
-            f"Save files to 'outputs/' directory:\n"
+            f"Save files to 'outputs/' directory using the FileWriterTool:\n"
             f"- {'index.html' if page_type == 'Single-Page' else 'Multiple HTML files'}"
         ),
         expected_output=(
@@ -747,7 +749,9 @@ chat_history = ChatHistory()
 def clean_outputs_directory():
     """Clean the outputs directory"""
     if os.path.exists("outputs"):
-        for file in glob.glob("outputs/*"):
+        if os.path.exists("outputs/project.zip"):
+            os.remove("outputs/project.zip")
+        for file in glob.glob("outputs/*.html"):
             try:
                 os.remove(file)
             except Exception as e:
@@ -756,17 +760,22 @@ def clean_outputs_directory():
 def create_zip_archive():
     """Create a ZIP archive of output files"""
     try:
-        html_files = glob.glob("outputs/*.html")
-        if len(html_files) > 1:
-            shutil.make_archive("outputs/project", 'zip', "outputs")
-            return "outputs/project.zip"
-        return None
+        # We need to archive the contents of 'outputs', not the directory itself
+        # This makes the zip file cleaner when extracted.
+        base_name = "outputs/project"
+        root_dir = "outputs"
+        # We'll temporarily remove the zip file if it exists to avoid adding it to itself
+        if os.path.exists(base_name + '.zip'):
+            os.remove(base_name + '.zip')
+            
+        shutil.make_archive(base_name, 'zip', root_dir, base_dir=None)
+        return "outputs/project.zip"
     except Exception as e:
         print(f"Error creating ZIP: {e}")
         return None
 
 # --- MAIN GENERATION FUNCTION ---
-def generate_website(website_type, project_name_input, name_method, page_type, custom_requirements, is_modification=False):
+def generate_website(website_type, project_name_input, name_method, page_type, custom_requirements):
     try:
         if llm is None:
             return (
@@ -781,13 +790,11 @@ def generate_website(website_type, project_name_input, name_method, page_type, c
             # Generate project name with AI
             project_name = f"{WEBSITE_TYPES[website_type]['emoji']} {website_type.split('/')[0]} Pro"
         
-        # Store in chat history
-        if not is_modification:
-            chat_history.reset()
-            chat_history.website_type = website_type
-            chat_history.project_name = project_name
-        
-        chat_history.add_message("User", f"Generate {website_type} - {page_type}")
+        # Reset chat history for a new project
+        chat_history.reset()
+        chat_history.website_type = website_type
+        chat_history.project_name = project_name
+        chat_history.add_message("User", f"Generate a new {website_type} website named '{project_name}' ({page_type}).")
         
         # Clean outputs
         clean_outputs_directory()
@@ -813,7 +820,7 @@ def generate_website(website_type, project_name_input, name_method, page_type, c
         html_files = glob.glob("outputs/*.html")
         if not html_files:
             return (
-                "❌ No files generated. Please try again.",
+                "❌ No HTML files were generated. The AI crew might have failed. Please check the console logs and try again.",
                 None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
             )
         
@@ -829,7 +836,9 @@ def generate_website(website_type, project_name_input, name_method, page_type, c
         success_message = f"✅ **{project_name}** generated successfully!\n"
         success_message += f"Type: {website_type} | Pages: {len(html_files)}"
         
-        chat_history.add_message("System", f"Generated {project_name}")
+        # **** FIX: Set current_project to enable modification ****
+        chat_history.current_project = project_name
+        chat_history.add_message("System", f"Successfully generated {project_name}. Files are available.")
         
         if is_multipage:
             zip_path = create_zip_archive()
@@ -853,42 +862,90 @@ def generate_website(website_type, project_name_input, name_method, page_type, c
         error_trace = traceback.format_exc()
         print(error_trace)
         return (
-            f"❌ Error: {str(e)}",
+            f"❌ An unexpected error occurred during generation: {str(e)}",
             None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
         )
 
+# **** FIX: Complete rewrite of modify_website function ****
 def modify_website(modification_request):
-    """Modify the existing website based on user request"""
+    """Modify the existing website based on user request and update the UI."""
+    # This function now returns a 5-element tuple to update all relevant UI components
+    error_response = (
+        "❌ Please generate a website first before requesting modifications.",
+        None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    )
+
     if not chat_history.current_project:
-        return "❌ Please generate a website first before requesting modifications."
-    
+        return error_response
+
+    if not modification_request or not modification_request.strip():
+        return ("⚠️ Please enter a modification request.", None, gr.update(), gr.update(), gr.update())
+
     chat_history.add_message("User", modification_request)
-    
-    # Create modification task
-    modification_task = Task(
-        description=(
-            f"Modify the existing {chat_history.website_type} website.\n\n"
-            f"Previous context:\n{chat_history.get_context()}\n\n"
-            f"Modification request: {modification_request}\n\n"
-            "Update the HTML files accordingly."
-        ),
-        expected_output="Modified HTML files with requested changes",
-        agent=developer,
-        tools=[file_write_tool]
-    )
-    
-    crew = Crew(
-        agents=[developer],
-        tasks=[modification_task],
-        process=Process.sequential
-    )
-    
+
     try:
+        modification_task = Task(
+            description=(
+                f"You are modifying an existing '{chat_history.website_type}' website named '{chat_history.project_name}'.\n\n"
+                "First, you MUST use the FileReadTool to read the existing HTML file(s) in the 'outputs' directory to understand the current code structure. The main file is likely 'index.html'.\n\n"
+                f"The user's new modification request is: '{modification_request}'\n\n"
+                "Based on this request, you must update the HTML file(s) in the 'outputs' directory. "
+                "Use the FileWriterTool to overwrite the existing file(s) with the necessary changes. "
+                "Do NOT create new files unless the request explicitly asks for a new page. Overwrite the existing ones.\n"
+                "Maintain the existing design system, colors, fonts, Tailwind CSS classes, and overall structure unless specifically asked to change them."
+            ),
+            expected_output="The complete, updated HTML code written back to the original file(s) in the 'outputs' directory using the FileWriterTool.",
+            agent=developer,
+            tools=[file_read_tool, file_write_tool]
+        )
+
+        crew = Crew(
+            agents=[developer],
+            tasks=[modification_task],
+            process=Process.sequential,
+            verbose=True
+        )
+        
         result = crew.kickoff()
-        chat_history.add_message("System", f"Modified: {modification_request[:50]}...")
-        return "✅ Modifications applied successfully!"
+        chat_history.add_message("System", f"Modification applied: {modification_request[:50]}...")
+
+        # After modification, refresh the UI with the new files
+        html_files = glob.glob("outputs/*.html")
+        if not html_files:
+            return "❌ Modification failed: No HTML files found after update.", None, gr.update(), gr.update(), gr.update(visible=False)
+        
+        preview_file = "outputs/index.html" if os.path.exists("outputs/index.html") else html_files[0]
+        with open(preview_file, "r", encoding="utf-8") as file:
+            html_content = file.read()
+
+        preview_html = f'<iframe srcdoc="{html_content.replace(chr(34), "&quot;")}" width="100%" height="800px" style="border: 2px solid #e0e0e0; border-radius: 8px;"></iframe>'
+        
+        is_multipage = len(html_files) > 1
+        success_message = f"✅ Modifications applied successfully to **{chat_history.project_name}**!"
+        
+        if is_multipage:
+            zip_path = create_zip_archive()
+            return (
+                success_message,
+                preview_html,
+                gr.update(value=preview_file, visible=True),
+                gr.update(value=zip_path, visible=True) if zip_path else gr.update(visible=False),
+                gr.update(visible=True)
+            )
+        else:
+            return (
+                success_message,
+                preview_html,
+                gr.update(value=preview_file, visible=True),
+                gr.update(visible=False),
+                gr.update(visible=True)
+            )
+
     except Exception as e:
-        return f"❌ Error applying modifications: {str(e)}"
+        error_trace = traceback.format_exc()
+        print(error_trace)
+        return f"❌ Error applying modifications: {str(e)}", None, gr.update(), gr.update(), gr.update(visible=False)
+
 
 # --- GRADIO INTERFACE ---
 with gr.Blocks(theme=gr.themes.Soft(), css="""
@@ -1026,16 +1083,16 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
             gr.Markdown(
                 """
                 ### Make Changes to Your Generated Website
-                After generating a website, you can request modifications here.
+                After generating a website, you can request modifications here. Your changes will update the live preview.
                 """
             )
             modification_input = gr.Textbox(
                 lines=3,
-                placeholder="E.g., 'Change the hero background to blue gradient', 'Add more testimonials', 'Make the navigation sticky'",
+                placeholder="E.g., 'Change the hero background to a blue gradient', 'Add more testimonials', 'Make the navigation sticky'",
                 label="What would you like to change?"
             )
             modify_btn = gr.Button("🔧 Apply Modifications", variant="secondary")
-            modification_status = gr.Markdown()
+            modification_status = gr.Markdown() # This will now be the output for modification status
     
     # Download Section
     with gr.Row():
@@ -1101,10 +1158,12 @@ with gr.Blocks(theme=gr.themes.Soft(), css="""
         outputs=[status_output, preview_output, download_single, download_zip, view_code_btn]
     )
     
+    # **** FIX: Update click handler for the modify button ****
     modify_btn.click(
         fn=modify_website,
         inputs=[modification_input],
-        outputs=[modification_status]
+        # The outputs now include the preview and download components to be refreshed
+        outputs=[modification_status, preview_output, download_single, download_zip, view_code_btn]
     )
     
     view_code_btn.click(