Update app.py

app.py

@@ -31,15 +31,24 @@ prompt_list = [
 ]
 
 # Load model globally (CPU-safe to avoid startup CUDA errors)
-pipe = pipeline("text-generation", model="xingyaoww/CodeActAgent-Mistral-7b-v0.1", device_map=None, torch_dtype=torch.float16)  # float16 fallback for init; no 'auto'
+pipe = pipeline("text-generation", model="xingyaoww/CodeActAgent-Mistral-7b-v0.1", device_map=None, torch_dtype=torch.float16)  # float16 for init; no 'auto'
 
 # Generator function with GPU decorator
-@spaces.GPU(duration=180)  # 180s for safety with multi-turn/model move
+@spaces.GPU(duration=180)  # 180s for safety
 def run_agent(user_content):
-    yield "Allocating GPU... (may queue if busy)\n\n"
-    yield "Moving model to GPU and initializing...\n\n"
+    full_log = ""
+    current_code = ""
+    current_exec_output = ""
 
-    # Move to GPU here (unrestricted power move)
+    yield [current_code, current_exec_output, full_log]  # Initial empty
+
+    full_log += "Allocating GPU... (may queue if busy)\n\n"
+    yield [current_code, current_exec_output, full_log]
+
+    full_log += "Moving model to GPU and initializing...\n\n"
+    yield [current_code, current_exec_output, full_log]
+
+    # Move to GPU here
     device = torch.device('cuda')
     pipe.model.to(device)
     pipe.device = device
@@ -60,35 +69,42 @@ def run_agent(user_content):
     max_turns = 10
     turn = 0
 
-    yield f"### Starting simulation for prompt: '{user_content}'\n\n"
+    full_log += f"### Starting simulation for prompt: '{user_content}'\n\n"
+    yield [current_code, current_exec_output, full_log]
 
     while turn < max_turns:
         # Generate response
         result = pipe(messages, max_new_tokens=512)
         assistant_content = result[0]['generated_text'][-1]['content']
         
-        yield f"**Assistant (Turn {turn+1}):** {assistant_content}\n\n"
+        full_log += f"**Assistant (Turn {turn+1}):** {assistant_content}\n\n"
+        yield [current_code, current_exec_output, full_log]
         
         # Stop checks
         if re.search(r'(task complete|done|final answer)', assistant_content.lower()):
-            yield "Detected completion keyword. Stopping.\n"
+            full_log += "Detected completion keyword. Stopping.\n"
+            yield [current_code, current_exec_output, full_log]
             break
         
         # Extract <execute>
         execute_match = re.search(r'<execute>(.*?)</execute>', assistant_content, re.DOTALL)
         if not execute_match:
-            yield "No code to execute. Task likely complete.\n"
+            full_log += "No code to execute. Task likely complete.\n"
+            yield [current_code, current_exec_output, full_log]
             break
         
         code = execute_match.group(1).strip()
+        current_code = code
         
         # Loop detection
         if code in prev_codes:
-            yield "Repeated code detected. Possible infinite loop—stopping.\n"
+            full_log += "Repeated code detected. Possible infinite loop—stopping.\n"
+            yield [current_code, current_exec_output, full_log]
             break
         prev_codes.add(code)
         
-        yield f"**Executing code:**\n```\n{code}\n```\n\n"
+        full_log += f"**Executing code:**\n```\n{code}\n```\n\n"
+        yield [current_code, current_exec_output, full_log]
         
         # Exec with capture
         old_stdout = sys.stdout
@@ -101,11 +117,14 @@ def run_agent(user_content):
         finally:
             sys.stdout = old_stdout
         
-        yield f"**Execution Output:** {exec_output}\n\n"
+        current_exec_output = exec_output
+        full_log += f"**Execution Output:** {exec_output}\n\n"
+        yield [current_code, current_exec_output, full_log]
         
         # Success stop: If output is pure digit (index), assume done
         if re.match(r'^\d+$', exec_output.strip()):
-            yield "Pure index output detected. Task successful—stopping.\n"
+            full_log += "Pure index output detected. Task successful—stopping.\n"
+            yield [current_code, current_exec_output, full_log]
             break
         
         # Append feedback
@@ -118,26 +137,36 @@ def run_agent(user_content):
     if 'exec_output' in locals():
         final_index = re.search(r'(\d+)$', exec_output)
         if final_index:
-            yield f"**Extracted Result:** Index {final_index.group(1)}\n"
+            full_log += f"**Extracted Result:** Index {final_index.group(1)}\n"
+            yield [current_code, current_exec_output, full_log]
         else:
-            yield "No clear index found—check errors.\n"
+            full_log += "No clear index found—check errors.\n"
+            yield [current_code, current_exec_output, full_log]
     else:
-        yield "No execution output.\n"
+        full_log += "No execution output.\n"
+        yield [current_code, current_exec_output, full_log]
 
-    yield f"### End of simulation for prompt: '{user_content}'\n"
+    full_log += f"### End of simulation for prompt: '{user_content}'\n"
+    yield [current_code, current_exec_output, full_log]
 
 # Gradio interface
 with gr.Blocks(title="Code Agent Simulator") as demo:
     gr.Markdown("# Code Agent Simulator on Hugging Face Spaces\nEnter a coding task prompt, and watch the agent simulate execution in real-time.")
     
     input_prompt = gr.Textbox(label="Enter your prompt", placeholder="e.g., Implement binary search...")
-    output_log = gr.Textbox(value="", lines=30, autoscroll=True, show_label=True, label="Simulation Log")
+    
+    with gr.Row():
+        generated_code = gr.Code(label="Generated Code", language="python", lines=15, show_label=True)
+        exec_output = gr.Textbox(label="Execution Output", lines=15, show_label=True)
+    
+    full_log = gr.Textbox(label="Full Simulation Log", lines=20, autoscroll=True, show_label=True)
+    
     run_button = gr.Button("Run Simulation")
     
     examples = gr.Examples(examples=prompt_list, inputs=[input_prompt])
     
-    # On click, run the generator and stream to output
-    run_button.click(fn=run_agent, inputs=input_prompt, outputs=output_log)
+    # On click, run the generator and stream to multiple outputs
+    run_button.click(fn=run_agent, inputs=input_prompt, outputs=[generated_code, exec_output, full_log])
 
 # Launch (disable SSR for stability, enable debug for logs)
 if __name__ == "__main__":