Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 import os
 import re
+# import traceback # Optional: for more detailed error logging if needed
 
 # --- Configuration ---
 API_TOKEN = os.getenv("HF_TOKEN", None)
@@ -14,46 +15,54 @@ try:
 except Exception as e:
     raise gr.Error(f"Failed to initialize model client for {MODEL}. Error: {e}. Check HF_TOKEN and model availability.")
 
-# --- Helper Function to Parse and Format Code ---
-def parse_and_format_code(raw_response: str) -> str:
+
+# --- Helper Function to Parse Code into Files ---
+def parse_code_into_files(raw_response: str) -> dict:
     """
     Parses raw AI output containing .TAB separators
-    and formats it for display in a single code block.
+    into a dictionary where keys are filenames and values are code blocks.
+    Returns keys like 'index.html', 'backend_file', 'backend_filename', 'backend_language'.
     """
+    files = {}
     # Default filename for the first block if no TAB is present or before the first TAB
-    default_filename = "index.html"
-    separator_pattern = r'\.TAB\[NAME=([^\]]+)\]\n?'
-
-    filenames = re.findall(r'\.TAB\[NAME=([^\]]+)\]', raw_response)
-    code_blocks = re.split(separator_pattern, raw_response)
+    default_first_filename = "index.html"
+    separator_pattern = r'\.TAB\[NAME=([^\]]+)\]\n?' # Capture filename
 
-    formatted_output = []
+    # Find all separators and their positions
+    matches = list(re.finditer(separator_pattern, raw_response))
 
-    # Handle the first block (before any potential separator)
-    first_block = code_blocks[0].strip()
+    start_index = 0
+    # Handle the first file (always assume index.html for now)
+    first_separator_pos = matches[0].start() if matches else len(raw_response)
+    first_block = raw_response[start_index:first_separator_pos].strip()
     if first_block:
-        formatted_output.append(f"--- START FILE: {default_filename} ---\n\n{first_block}\n\n--- END FILE: {default_filename} ---")
+         files[default_first_filename] = first_block
+
+    # Handle the second file (if separator exists)
+    if matches:
+        backend_filename = matches[0].group(1).strip() # Get filename from first match
+        start_index = matches[0].end() # Start after the first separator
 
-    # Handle subsequent blocks associated with filenames found by the separator
-    # re.split with capturing groups includes the separator AND the filename capture group
-    # in the results. We need to iterate carefully.
-    # Example: ['code1', 'app.py', 'code2', 'style.css', 'code3']
-    idx = 1
-    while idx < len(code_blocks) -1 :
-        filename = code_blocks[idx] # This should be the filename captured by the pattern
-        code = code_blocks[idx + 1].strip() # This should be the code after the separator
-        if code : # Only add if there's actual code content
-             formatted_output.append(f"--- START FILE: {filename} ---\n\n{code}\n\n--- END FILE: {filename} ---")
-        idx += 2 # Move past the filename and the code block
+        # Find the position of the *next* separator, or end of string
+        second_separator_pos = matches[1].start() if len(matches) > 1 else len(raw_response)
+        backend_code = raw_response[start_index:second_separator_pos].strip()
 
-    # If no separators were found, filenames list will be empty, and only the first block runs.
-    # If separators were found, the loop processes the rest.
+        if backend_code:
+            files['backend_file'] = backend_code
+            files['backend_filename'] = backend_filename
+            # Determine language from filename extension
+            if backend_filename.endswith(".py"):
+                files['backend_language'] = 'python'
+            elif backend_filename.endswith(".js"):
+                files['backend_language'] = 'javascript'
+            elif backend_filename.endswith(".css"):
+                 files['backend_language'] = 'css'
+            else:
+                 files['backend_language'] = None # Default to plain text
 
-    if not formatted_output: # Handle case where input was empty or only whitespace
-        return raw_response # Return the original if parsing yields nothing
+    # If more files were generated (more separators), they are currently ignored by this simple parser.
 
-    # Join the formatted blocks with clear separation
-    return "\n\n\n".join(formatted_output)
+    return files
 
 
 # --- Core Code Generation Function ---
@@ -63,23 +72,26 @@ def generate_code(
     max_tokens: int,
     temperature: float,
     top_p: float,
+    progress=gr.Progress(track_ τότε=True) # Add progress tracker
 ):
     print(f"Generating code for: {prompt[:100]}... | Backend: {backend_choice}")
+    progress(0, desc="Initializing Request...")
 
+    # System message remains the same - instructing the AI on format
     system_message = (
-        "You are an AI that generates website code. You MUST ONLY output the raw code, without any conversational text like 'Here is the code' or explanations before or after the code blocks. "
-        "You MUST NOT wrap the code in markdown fences like ```html, ```python, or ```js. "
-        "If the user requests 'Static' or the prompt clearly implies only frontend code, generate ONLY the content for the `index.html` file. "
-        "If the user requests 'Flask' or 'Node.js' and the prompt requires backend logic, you MUST generate both the `index.html` content AND the corresponding main backend file content (e.g., `app.py` for Flask, `server.js` or `app.js` for Node.js). "
-        "When generating multiple files, you MUST separate them EXACTLY as follows: "
-        "1. Output the complete code for the first file (e.g., `index.html`). "
-        "2. On a new line immediately after the first file's code, add the separator '.TAB[NAME=filename.ext]' (e.g., '.TAB[NAME=app.py]' or '.TAB[NAME=server.js]'). "
-        "3. On the next line, immediately start the code for the second file. "
-        "Generate only the necessary files (usually index.html and potentially one backend file). "
-        "The generated website code must be SFW and have minimal errors. "
-        "Only include comments where user modification is strictly required. Avoid explanatory comments. "
-        "If the user asks you to create code that is NOT for a website, you MUST respond ONLY with the exact phrase: "
-        "'hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-('"
+         "You are an AI that generates website code. You MUST ONLY output the raw code, without any conversational text like 'Here is the code' or explanations before or after the code blocks. "
+         "You MUST NOT wrap the code in markdown fences like ```html, ```python, or ```js. "
+         "If the user requests 'Static' or the prompt clearly implies only frontend code, generate ONLY the content for the `index.html` file. "
+         "If the user requests 'Flask' or 'Node.js' and the prompt requires backend logic, you MUST generate both the `index.html` content AND the corresponding main backend file content (e.g., `app.py` for Flask, `server.js` or `app.js` for Node.js). "
+         "When generating multiple files, you MUST separate them EXACTLY as follows: "
+         "1. Output the complete code for the first file (e.g., `index.html`). "
+         "2. On a new line immediately after the first file's code, add the separator '.TAB[NAME=filename.ext]' (e.g., '.TAB[NAME=app.py]' or '.TAB[NAME=server.js]'). "
+         "3. On the next line, immediately start the code for the second file. "
+         "Generate only the necessary files (usually index.html and potentially one backend file). "
+         "The generated website code must be SFW and have minimal errors. "
+         "Only include comments where user modification is strictly required. Avoid explanatory comments. "
+         "If the user asks you to create code that is NOT for a website, you MUST respond ONLY with the exact phrase: "
+         "'hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-('"
     )
 
     user_prompt = f"USER_PROMPT = {prompt}\nUSER_BACKEND = {backend_choice}"
@@ -89,10 +101,12 @@ def generate_code(
         {"role": "user", "content": user_prompt}
     ]
 
-    response_stream = ""
     full_response = ""
+    token_count = 0
+    est_total_tokens = max_tokens # Rough estimate for progress
 
     try:
+        progress(0.1, desc="Sending Request to Model...")
         stream = client.chat_completion(
             messages=messages,
             max_tokens=max_tokens,
@@ -100,27 +114,31 @@ def generate_code(
             temperature=temperature,
             top_p=top_p,
         )
+
+        progress(0.2, desc="Receiving Stream...")
         for message in stream:
             token = message.choices[0].delta.content
             if isinstance(token, str):
-                response_stream += token
                 full_response += token
-                # Yield intermediate stream for responsiveness during generation
-                yield response_stream # Show raw output as it comes
+                token_count += 1
+                # Update progress based on tokens received vs max_tokens
+                # Adjust the scaling factor (e.g., 0.7) as needed
+                prog = min(0.2 + (token_count / est_total_tokens) * 0.7, 0.9)
+                progress(prog, desc="Generating Code...")
 
-        # --- Post-processing (After Stream Ends) ---
-        cleaned_response = full_response.strip()
 
+        progress(0.9, desc="Processing Response...")
+        # --- Post-processing ---
+        cleaned_response = full_response.strip()
         # Fallback fence removal
         cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", cleaned_response)
         cleaned_response = re.sub(r"\n?\s*```\s*$", "", cleaned_response)
         # Remove potential chat markers
         cleaned_response = re.sub(r"<\s*\|?\s*(user|system|assistant)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE).strip()
         # Remove common conversational phrases (if they slip through)
-        common_phrases = [
+        common_phrases = [ # Simplified list as prompt should handle most
             "Here is the code:", "Okay, here is the code:", "Here's the code:",
-            "Sure, here is the code you requested:", "Let me know if you need anything else.",
-            "```html", "```python", "```javascript", "```",
+            "Sure, here is the code you requested:",
         ]
         temp_response_lower = cleaned_response.lower()
         for phrase in common_phrases:
@@ -131,19 +149,39 @@ def generate_code(
         # Check for refusal message
         refusal_message = "hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-("
         if refusal_message in full_response:
-             yield refusal_message # Yield the exact refusal message
-             return # Stop processing
+             # Return updates to clear both code blocks and show refusal in the first
+             progress(1, desc="Refusal Message Generated")
+             return gr.update(value=refusal_message, language=None, visible=True), gr.update(value="", visible=False)
 
-        # --- PARSE and FORMAT the final cleaned response ---
-        formatted_final_code = parse_and_format_code(cleaned_response)
+        # --- PARSE the final cleaned response into files ---
+        parsed_files = parse_code_into_files(cleaned_response)
 
-        # Yield the final, formatted version replacing the streamed content
-        yield formatted_final_code
+        html_code = parsed_files.get("index.html", "")
+        backend_code = parsed_files.get("backend_file", "")
+        backend_filename = parsed_files.get("backend_filename", "Backend")
+        backend_language = parsed_files.get("backend_language", None)
+
+        # --- Prepare Gradio Updates ---
+        # Update for the HTML code block (always visible)
+        html_update = gr.update(value=html_code, language='html', visible=True)
+
+        # Update for the Backend code block (visible only if backend code exists)
+        if backend_code:
+            backend_update = gr.update(value=backend_code, language=backend_language, label=backend_filename, visible=True)
+        else:
+            backend_update = gr.update(value="", visible=False) # Hide if no backend code
+
+        progress(1, desc="Done")
+        # Return tuple of updates for the outputs list
+        return html_update, backend_update
 
     except Exception as e:
         print(f"ERROR during code generation: {e}") # Log detailed error
-        # traceback.print_exc() # Uncomment for full traceback in logs
-        yield f"## Error\n\nFailed to generate or process code.\n**Reason:** {e}"
+        # traceback.print_exc() # Uncomment for full traceback
+        progress(1, desc="Error Occurred")
+        error_message = f"## Error\n\nFailed to generate or process code.\n**Reason:** {e}"
+        # Return updates to show error in the first block and hide the second
+        return gr.update(value=error_message, language=None, visible=True), gr.update(value="", visible=False)
 
 
 # --- Build Gradio Interface ---
@@ -151,22 +189,14 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
     gr.Markdown("# ✨ Website Code Generator ✨")
     gr.Markdown(
         "Describe the website you want. The AI will generate the necessary code.\n"
-        "If multiple files are generated (e.g., for Flask/Node.js), they will be shown below separated by `--- START FILE: filename ---` markers.\n" # Updated description
-        "**Output Format:**\n"
-        "- No explanations, just code.\n"
-        "- Multiple files separated by file markers.\n" # Updated description
-        "- Minimal necessary comments only.\n\n"
-        "**Rules:**\n"
-        "- Backend choice guides the AI on whether to include server-side code.\n"
-        "- Always SFW and aims for minimal errors.\n"
-        "- Only generates website-related code."
+        "If multiple files are generated (e.g., for Flask/Node.js), they will appear in separate tabs below." # Updated description
     )
 
     with gr.Row():
         with gr.Column(scale=2):
             prompt_input = gr.Textbox(
                 label="Website Description",
-                placeholder="e.g., A Flask app with a form that stores data in a variable.",
+                placeholder="e.g., A Flask app with a simple chat using Socket.IO",
                 lines=6,
             )
             backend_radio = gr.Radio(
@@ -178,12 +208,29 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
             generate_button = gr.Button("✨ Generate Website Code", variant="primary")
 
         with gr.Column(scale=3):
-            code_output = gr.Code(
-                label="Generated Code (Scroll for multiple files)", # Updated label
-                language=None, # Keep as None for mixed/plain text display
-                lines=30,
-                interactive=False, # Keep non-interactive
-            )
+            # Define Tabs to hold the code outputs
+            with gr.Tabs(elem_id="code-tabs") as code_tabs:
+                # Tab 1: Always present for HTML
+                with gr.Tab("index.html", elem_id="html-tab") as html_tab:
+                     html_code_output = gr.Code(
+                         label="index.html", # Label for the code block itself
+                         language="html",
+                         lines=25, # Adjusted lines slightly
+                         interactive=False,
+                         elem_id="html_code", # Unique ID for targeting
+                     )
+                # Tab 2: For Backend code, initially hidden
+                with gr.Tab("Backend", elem_id="backend-tab", visible=False) as backend_tab:
+                     backend_code_output = gr.Code(
+                         label="Backend Code", # Label will be updated dynamically
+                         language=None, # Language updated dynamically
+                         lines=25,
+                         interactive=False,
+                         elem_id="backend_code", # Unique ID for targeting
+                         visible=False # Component also starts hidden
+                     )
+                # Add more tabs here if needed (e.g., for CSS) following the same pattern
+
 
     with gr.Accordion("Advanced Settings", open=False):
         max_tokens_slider = gr.Slider(
@@ -196,15 +243,15 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
             minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P"
         )
 
-    # The click function now yields the final formatted string AFTER streaming
+    # The click function now targets the specific code blocks within the tabs
     generate_button.click(
         fn=generate_code,
         inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
-        outputs=code_output,
+        # The outputs list MUST match the order and number of code blocks we want to update
+        outputs=[html_code_output, backend_code_output],
     )
 
 if __name__ == "__main__":
     if not API_TOKEN:
         print("Warning: HF_TOKEN environment variable not set. Using anonymous access.")
-    # Increased max_size slightly for potentially larger combined outputs
-    demo.queue(max_size=15).launch()
+    demo.queue(max_size=10).launch() # Allow queueing