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 debugging
 
 API_TOKEN = os.getenv("HF_TOKEN", None)
 MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
@@ -12,52 +13,64 @@ try:
 except Exception as e:
     raise gr.Error(f"Failed to initialize model client for {MODEL}. Error: {e}. Check HF_TOKEN and model availability.")
 
-def parse_code_into_files(raw_response: str) -> dict:
-    files = {}
-    default_first_filename = "index.html"
+# --- Helper Function to Parse Code during Streaming ---
+def parse_streaming_code(current_response: str) -> dict:
+    """
+    Parses potentially incomplete AI output stream.
+    Identifies if .TAB separator is present and splits code accordingly.
+    Returns dict with html_code, backend_code, filename, language, and visibility flag.
+    """
+    files = {
+        'html_code': '',
+        'backend_code': '',
+        'backend_filename': 'Backend', # Default label
+        'backend_language': None,
+        'backend_visible': False       # Default visibility
+    }
     separator_pattern = r'\.TAB\[NAME=([^\]]+)\]\n?'
-
-    matches = list(re.finditer(separator_pattern, raw_response))
-
-    start_index = 0
-    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:
-         files[default_first_filename] = first_block
-
-    if matches:
-        backend_filename = matches[0].group(1).strip()
-        start_index = matches[0].end()
-
-        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 backend_code:
-            files['backend_file'] = backend_code
-            files['backend_filename'] = backend_filename
-            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
+    match = re.search(separator_pattern, current_response)
+
+    if match:
+        # Separator found in the stream so far
+        html_part = current_response[:match.start()].strip()
+        backend_part = current_response[match.end():].strip() # Code after separator
+        backend_filename = match.group(1).strip()
+
+        files['html_code'] = html_part
+        files['backend_code'] = backend_part
+        files['backend_filename'] = backend_filename
+        files['backend_visible'] = True # Make backend visible
+
+        # Determine language
+        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
+    else:
+        # No separator found yet, assume all content is HTML
+        files['html_code'] = current_response.strip()
+        # Keep backend_visible as False
 
     return files
 
+# --- Minimal Cleaning for Intermediate Stream Chunks ---
+def clean_intermediate_stream(text: str) -> str:
+    """ Basic cleaning for streaming chunks (e.g., remove chat markers). """
+    cleaned = re.sub(r"<\s*\|?\s*(user|system|assistant)\s*\|?\s*>", "", text, flags=re.IGNORECASE)
+    # Avoid stripping aggressively during stream as it might remove partial code
+    return cleaned
+
+# --- Core Code Generation Function - Modified for Streaming UI Updates ---
 def generate_code(
     prompt: str,
     backend_choice: str,
     max_tokens: int,
     temperature: float,
     top_p: float,
-    progress=gr.Progress(track_tqdm=True)
 ):
-    print(f"Generating code for: {prompt[:100]}... | Backend: {backend_choice}")
-    progress(0, desc="Initializing Request...")
+    print(f"Streaming code generation for: {prompt[:100]}... | Backend: {backend_choice}")
 
-    system_message = (
+    system_message = ( # System message remains the same
          "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. "
@@ -81,11 +94,17 @@ def generate_code(
     ]
 
     full_response = ""
-    token_count = 0
-    est_total_tokens = max_tokens
+    # Initialize state for UI updates
+    current_html = ""
+    current_backend = ""
+    current_backend_label = "Backend"
+    current_backend_lang = None
+    is_backend_visible = False
+
+    # Initial clear of outputs
+    yield gr.update(value="", visible=True), gr.update(visible=False), gr.update(value="", visible=False)
 
     try:
-        progress(0.1, desc="Sending Request to Model...")
         stream = client.chat_completion(
             messages=messages,
             max_tokens=max_tokens,
@@ -94,65 +113,86 @@ def generate_code(
             top_p=top_p,
         )
 
-        progress(0.2, desc="Receiving Stream...")
         for message in stream:
             token = message.choices[0].delta.content
             if isinstance(token, str):
                 full_response += token
-                token_count += 1
-                prog = min(0.2 + (token_count / est_total_tokens) * 0.7, 0.9)
-                progress(prog, desc="Generating Code...")
-
-
-        progress(0.9, desc="Processing Response...")
-        cleaned_response = full_response.strip()
-        cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", cleaned_response)
-        cleaned_response = re.sub(r"\n?\s*```\s*$", "", cleaned_response)
-        cleaned_response = re.sub(r"<\s*\|?\s*(user|system|assistant)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE).strip()
-        common_phrases = [
-            "Here is the code:", "Okay, here is the code:", "Here's the code:",
-            "Sure, here is the code you requested:",
-        ]
-        temp_response_lower = cleaned_response.lower()
+                # Clean intermediate response minimally
+                cleaned_response_chunk = clean_intermediate_stream(full_response)
+
+                # Parse the *entire accumulated* cleaned response on each iteration
+                parsed_state = parse_streaming_code(cleaned_response_chunk)
+
+                # Update state variables
+                current_html = parsed_state['html_code']
+                current_backend = parsed_state['backend_code']
+                current_backend_label = parsed_state['backend_filename']
+                current_backend_lang = parsed_state['backend_language']
+                is_backend_visible = parsed_state['backend_visible'] # This determines visibility
+
+                # Prepare Gradio updates based on the *current* parsed state
+                html_update = gr.update(value=current_html)
+                # Update the backend tab's visibility
+                tab_update = gr.update(visible=is_backend_visible)
+                # Update the backend code block's content, label, language, and visibility
+                backend_code_update = gr.update(
+                    value=current_backend,
+                    label=current_backend_label,
+                    language=current_backend_lang,
+                    visible=is_backend_visible # Make code block visible *if* tab is visible
+                )
+
+                # Yield updates for html_code, backend_tab, backend_code
+                yield html_update, tab_update, backend_code_update
+
+        # --- Final Clean and Update after Stream ---
+        # Ensure the final state is clean and fully parsed
+        final_cleaned_response = full_response.strip()
+        # Remove fences/phrases missed during stream (optional but good practice)
+        final_cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", final_cleaned_response)
+        final_cleaned_response = re.sub(r"\n?\s*```\s*$", "", final_cleaned_response)
+        common_phrases = ["Here is the code:", "Okay, here is the code:", "Here's the code:", "Sure, here is the code you requested:"]
+        temp_lower = final_cleaned_response.lower()
         for phrase in common_phrases:
-            if temp_response_lower.startswith(phrase.lower()):
-                cleaned_response = cleaned_response[len(phrase):].lstrip()
-                temp_response_lower = cleaned_response.lower()
+            if temp_lower.startswith(phrase.lower()):
+                final_cleaned_response = final_cleaned_response[len(phrase):].lstrip()
+                temp_lower = final_cleaned_response.lower()
 
+        # Check for refusal message in the final response
         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:
-             progress(1, desc="Refusal Message Generated")
-             return gr.update(value=refusal_message, language=None, visible=True), gr.update(value="", visible=False, label="Backend")
-
-        parsed_files = parse_code_into_files(cleaned_response)
-
-        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)
-
-        html_update = gr.update(value=html_code, language='html', visible=True)
-
-        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, label="Backend")
+        if refusal_message in final_cleaned_response:
+             yield gr.update(value=refusal_message), gr.update(visible=False), gr.update(value="", visible=False)
+             return # Stop processing
+
+        # Final parse
+        final_parsed_state = parse_streaming_code(final_cleaned_response)
+
+        # Final updates to ensure everything is correct
+        final_html_update = gr.update(value=final_parsed_state['html_code'])
+        final_tab_update = gr.update(visible=final_parsed_state['backend_visible'])
+        final_backend_code_update = gr.update(
+            value=final_parsed_state['backend_code'],
+            label=final_parsed_state['backend_filename'],
+            language=final_parsed_state['backend_language'],
+            visible=final_parsed_state['backend_visible']
+        )
+        yield final_html_update, final_tab_update, final_backend_code_update
 
-        progress(1, desc="Done")
-        return html_update, backend_update
 
     except Exception as e:
-        print(f"ERROR during code generation: {e}")
-        progress(1, desc="Error Occurred")
-        error_message = f"## Error\n\nFailed to generate or process code.\n**Reason:** {e}"
-        return gr.update(value=error_message, language=None, visible=True), gr.update(value="", visible=False, label="Backend")
+        print(f"ERROR during code generation stream: {e}")
+        # traceback.print_exc() # Uncomment for detailed traceback
+        error_message = f"## Error\n\nFailed during streaming.\n**Reason:** {e}"
+        # Show error in HTML block, hide backend tab and code
+        yield gr.update(value=error_message), gr.update(visible=False), gr.update(value="", visible=False)
 
 
+# --- Build Gradio Interface ---
 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 appear in separate tabs below."
+        "Describe the website you want. See code generated live.\n"
+        "If backend code is generated, a second tab will appear."
     )
 
     with gr.Row():
@@ -171,7 +211,9 @@ 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):
-            with gr.Tabs(elem_id="code-tabs") as code_tabs:
+            # Define Tabs structure
+            with gr.Tabs(elem_id="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",
@@ -180,14 +222,15 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
                          interactive=False,
                          elem_id="html_code",
                      )
+                # Tab 2: Backend - defined but starts hidden
                 with gr.Tab("Backend", elem_id="backend-tab", visible=False) as backend_tab:
                      backend_code_output = gr.Code(
-                         label="Backend",
-                         language=None,
+                         label="Backend", # Label updated dynamically if tab becomes visible
+                         language=None,   # Language updated dynamically
                          lines=25,
                          interactive=False,
                          elem_id="backend_code",
-                         visible=False
+                         visible=False # Code block also starts hidden
                      )
 
     with gr.Accordion("Advanced Settings", open=False):
@@ -201,10 +244,12 @@ 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"
         )
 
+    # Click function now targets html_code_output, backend_tab, and backend_code_output
     generate_button.click(
         fn=generate_code,
         inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
-        outputs=[html_code_output, backend_code_output],
+        # Outputs MUST match the number and order of updates yielded by the function
+        outputs=[html_code_output, backend_tab, backend_code_output],
     )
 
 if __name__ == "__main__":