Update app.py

app.py

@@ -8,7 +8,7 @@ def respond(message, history, system_message, max_tokens, temperature, top_p, ba
 You are a code-generation AI. You MUST generate a full website including an index.html file.
 Use only the {backend} backend structure. 
 Respond ONLY with raw code and file/folder structure. Do NOT explain or add commentary.
-    """.strip()
+""".strip()
 
     system_message = forced_system + "\n\n" + system_message
 
@@ -21,14 +21,14 @@ Respond ONLY with raw code and file/folder structure. Do NOT explain or add comm
     messages.append({"role": "user", "content": message})
 
     response = ""
-    for message in client.chat_completion(
+    for chunk in client.chat_completion(
         messages,
         max_tokens=max_tokens,
         stream=True,
         temperature=temperature,
         top_p=top_p,
     ):
-        token = message.choices[0].delta.content
+        token = chunk.choices[0].delta.content
         if token:
             response += token
             yield response
@@ -37,32 +37,43 @@ with gr.Blocks() as demo:
     gr.Markdown("# WebGen AI\nGenerate a complete website with your selected backend.")
 
     with gr.Row():
-        system_msg = gr.Textbox(value="You're a helpful assistant.", label="System Message", lines=2)
+        system_msg = gr.Textbox(value="You are a helpful assistant.", label="System Message")
         backend = gr.Dropdown(["Flask", "Static", "Node.js"], value="Static", label="Backend")
-    
+
     with gr.Row():
         max_tokens = gr.Slider(1, 2048, value=512, label="Max Tokens")
         temperature = gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature")
         top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
 
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Ask to generate a website...", label="Your Message")
+    user_input = gr.Textbox(label="Your Prompt", placeholder="Ask the AI to generate a website...")
+
+    history = []
 
-    def user_submit(message, chat_history):
-        return "", chat_history + [[message, None]]
+    def chat_submit(message):
+        nonlocal history
+        history.append([message, None])
+        return "", history
 
-    submit = gr.Button("Send")
+    send_btn = gr.Button("Send")
 
-    state = gr.State([])
+    def run_response(message, system_msg, max_tokens, temperature, top_p, backend):
+        nonlocal history
+        response_generator = respond(message, history, system_msg, max_tokens, temperature, top_p, backend)
+        final_response = ""
+        for chunk in response_generator:
+            final_response = chunk
+            yield history[:-1] + [[message, chunk]]
+        history[-1][1] = final_response
 
-    submit.click(
-        user_submit,
-        [msg, state],
-        [msg, state]
+    send_btn.click(
+        chat_submit,
+        inputs=[user_input],
+        outputs=[user_input, chatbot]
     ).then(
-        respond,
-        [msg, state, system_msg, max_tokens, temperature, top_p, backend],
-        chatbot
+        run_response,
+        inputs=[user_input, system_msg, max_tokens, temperature, top_p, backend],
+        outputs=chatbot
     )
 
 if __name__ == "__main__":