Update app.py

app.py

@@ -26,7 +26,6 @@ def generate_code(prompt: str, backend_choice: str, max_tokens: int, temperature
         {"role": "system", "content": system_message},
         {"role": "user", "content": user_prompt}
     ]
-    response_stream = ""
     full_response = ""
     try:
         stream = client.chat_completion(
@@ -39,9 +38,7 @@ def generate_code(prompt: str, backend_choice: str, max_tokens: int, temperature
         for message in stream:
             token = message.choices[0].delta.content
             if isinstance(token, str):
-                response_stream += token
                 full_response += token
-                yield {"raw": response_stream}
         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)
@@ -53,26 +50,22 @@ def generate_code(prompt: str, backend_choice: str, max_tokens: int, temperature
         for phrase in common_phrases:
             if cleaned_response.lower().startswith(phrase.lower()):
                 cleaned_response = cleaned_response[len(phrase):].lstrip()
-        yield {"final": cleaned_response.strip()}
+        return cleaned_response
     except Exception as e:
-        yield {"error": f"## Error\n\nFailed to generate code.\n**Reason:** {e}"}
+        return f"## Error\n\nFailed to generate code.\n**Reason:** {e}"
 
-def display_output(output_dict):
-    if "error" in output_dict:
-        return gr.Markdown(output_dict["error"])
-    if "raw" in output_dict:
-        return gr.Markdown("Generating...")
-    cleaned_text = output_dict["final"]
-    file_splits = re.split(r'TAB\.NAME=\{(.+?)\}', cleaned_text)
-    components = []
-    if not file_splits[0].strip():
-        file_splits = file_splits[1:]
-    for i in range(0, len(file_splits), 2):
-        filename = file_splits[i].strip()
-        content = file_splits[i+1].strip() if (i+1) < len(file_splits) else ""
-        components.append(gr.Markdown(f"### {filename}"))
-        components.append(gr.Code(value=content, label=filename, language="html" if filename.endswith(".html") else "python" if filename.endswith(".py") else "javascript" if filename.endswith(".js") else "text", interactive=False))
-    return components
+def split_files(full_code_text):
+    file_blocks = []
+    splits = re.split(r'TAB\.NAME=\{(.+?)\}', full_code_text)
+    if not splits[0].strip():
+        splits = splits[1:]
+    for i in range(0, len(splits), 2):
+        if i+1 >= len(splits):
+            continue
+        filename = splits[i].strip()
+        content = splits[i+1].strip()
+        file_blocks.append((filename, content))
+    return file_blocks
 
 with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
     gr.Markdown("# ✨ Website Code Generator ✨")
@@ -98,7 +91,14 @@ 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):
-            output_display = gr.Group()
+            main_output = gr.Code(
+                label="Full Generated Code",
+                language="html",
+                lines=25,
+                interactive=False,
+            )
+            extra_outputs = gr.Column()
+
     with gr.Accordion("Advanced Settings", open=False):
         max_tokens_slider = gr.Slider(
             minimum=512,
@@ -108,7 +108,8 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
             label="Max New Tokens"
         )
         temperature_slider = gr.Slider(
-            minimum=0.1, maximum=1.2,
+            minimum=0.1,
+            maximum=1.2,
             value=0.7,
             step=0.1,
             label="Temperature"
@@ -120,14 +121,20 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
             step=0.05,
             label="Top-P"
         )
+
+    def generate_and_display(prompt, backend, max_tokens, temperature, top_p):
+        full_code = generate_code(prompt, backend, max_tokens, temperature, top_p)
+        files = split_files(full_code)
+        file_outputs = []
+        for filename, content in files:
+            file_outputs.append(gr.Markdown.update(value=f"### {filename}"))
+            file_outputs.append(gr.Code.update(value=content, language="html" if filename.endswith(".html") else "python" if filename.endswith(".py") else "javascript" if filename.endswith(".js") else "text"))
+        return full_code, file_outputs
+
     generate_button.click(
-        fn=generate_code,
+        fn=generate_and_display,
         inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
-        outputs=output_display,
-    ).then(
-        fn=display_output,
-        inputs=output_display,
-        outputs=output_display,
+        outputs=[main_output, extra_outputs]
     )
 
 if __name__ == "__main__":