Update app.py

app.py

@@ -2,7 +2,7 @@ from smolagents import CodeAgent, DuckDuckGoSearchTool, OpenAIServerModel, tool,
 import gradio as gr
 import requests
 import pytz
-from typing import Optional, Tuple, Union, Any, List, Dict, Iterator
+from typing import Optional, Tuple, Union, Any, List, Dict
 import re
 # from google.colab import userdata # Assuming Colab environment
 import io
@@ -273,96 +273,36 @@ def agent_thread_func(agent, query, log_queue, result_queue):
 
 # Generator function for Gradio streaming
 # REMOVED the return type hint -> Iterator[...]
-def run_agent_wrapper(query: str):
+# --- REPLACE the current run_agent_wrapper function WITH THIS ---
+def run_agent_wrapper(query: str) -> Tuple[str, str]:
     """
-    Runs the agent in a thread, captures stdout via a queue, and yields updates
-    for Gradio streaming. Uses Chatbot 'messages' format for reasoning.
-    Yields: (chatbot_history, final_answer_status)
+    Runs the height_agent synchronously and captures its stdout (reasoning steps).
+    Returns (reasoning_log, final_answer). NO STREAMING.
     """
     if height_agent is None:
         error_msg = initialization_error_message or "Agent not initialized."
-        # Yield error in the 'messages' format
-        yield ([{"role": "assistant", "content": error_msg}], "Error: Agent not initialized.")
-        return # Stop the generator
-
-    log_queue = queue.Queue()
-    result_queue = queue.Queue()
-    # History will be a list of dictionaries: [{"role": "assistant", "content": "..."}]
-    # We'll just use one dictionary and update its content for the streaming log
-    chatbot_history = []
-    current_log_message = "" # Accumulate lines into one message block
-    final_answer = "⏳ Running..." # Initial status
-
-    # Initial yield to clear previous state and show "Running"
-    # Yield empty history initially, or a starting message
-    yield ([], final_answer)
-
-    # Start the agent thread
-    thread = threading.Thread(
-        target=agent_thread_func,
-        args=(height_agent, query, log_queue, result_queue)
-    )
-    thread.start()
+        return (error_msg, "Agent failed to initialize. See reasoning log.")
 
-    while True:
-        try:
-            # Check for new log messages (non-blocking)
-            log_line = log_queue.get_nowait()
-
-            if log_line is None: # End-of-logs signal
-                break
-
-            # Append new line to the current log message block
-            # Add line breaks for readability in the chatbot
-            current_log_message += log_line + "\n"
-
-            # Update the chatbot history: Replace the last message or add a new one
-            # Simplified: Update a single assistant message with the whole log
-            chatbot_history = [{"role": "assistant", "content": current_log_message}]
-
-            yield (chatbot_history, final_answer) # Yield updated log
-
-        except queue.Empty:
-            # No new message, brief pause to prevent busy-waiting
-            # Also check if the thread is still alive; if not, break (error case)
-            if not thread.is_alive() and result_queue.empty():
-                print("Warning: Agent thread finished unexpectedly without result.")
-                # Attempt to retrieve any remaining logs
-                try:
-                    while True: # Get all remaining logs
-                        log_line = log_queue.get_nowait()
-                        if log_line: current_log_message += log_line + "\n"
-                        else: break # Should not happen if None was already processed, but safety
-                except queue.Empty:
-                    pass # No more logs
-                current_log_message += "\nError: Agent stopped unexpectedly."
-                chatbot_history = [{"role": "assistant", "content": current_log_message}]
-                final_answer = "Error: Agent stopped unexpectedly."
-                yield (chatbot_history, final_answer)
-                return # Stop
-
-            time.sleep(0.1) # Pause briefly
-
-    # Agent thread has finished (log_queue received None)
-    thread.join() # Wait for the thread to fully terminate
-
-    # Get the final result or exception
-    final_result = result_queue.get()
-
-    if isinstance(final_result, Exception):
-        final_answer = f"Error during execution: {final_result}"
-        # Append error to the chatbot log
-        error_log = f"\n\n--- EXECUTION ERROR ---\n{final_result}"
-        current_log_message += error_log
-        chatbot_history = [{"role": "assistant", "content": current_log_message}]
-    else:
-        final_answer = final_result # This is the actual final answer string
-
-    # Final yield with the complete log and the final answer
-    # Ensure history is in the correct format before the final yield
-    chatbot_history = [{"role": "assistant", "content": current_log_message}]
-    yield (chatbot_history, final_answer)
+    print(f"\n--- Running agent for query: '{query}' ---")
+    log_stream = io.StringIO()
+    final_answer = "Error during execution." # Default message
+
+    try:
+        # Redirect stdout to capture prints from agent.run()
+        with contextlib.redirect_stdout(log_stream):
+            # Run the agent directly (prints are captured)
+            final_answer = height_agent.run(query)
+            print("\n--- Agent execution finished. ---") # Add marker to log
+    except Exception as e:
+        print(f"\n--- Error during agent execution wrapper: {e} ---")
+        traceback.print_exc(file=log_stream) # Print exception to log stream
+        final_answer = f"An error occurred in the wrapper. See reasoning log. Error: {e}"
+    finally:
+        reasoning_log = log_stream.getvalue()
+        log_stream.close()
+        print("--- Finished capturing stdout. ---") # Log to console, not captured
 
+    return reasoning_log, final_answer
 
 # --- Build Gradio Interface Manually with gr.Blocks ---
 print("--- Building Gradio Interface with gr.Blocks ---")
@@ -386,18 +326,21 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
     gr.Markdown("## Agent Reasoning Steps")
     # --- CHANGE THIS ---
     # reasoning_output = gr.Code(label="Reasoning Log", language="markdown", interactive=False, lines=20)
-    reasoning_output_chatbot = gr.Chatbot(
+    # --- REPLACE Chatbot definition WITH THIS ---
+    reasoning_output = gr.Code(
     label="Reasoning Log",
-    height=500,
-    type="messages" # <<< ADD THIS
-    )
+    language="markdown", # Use markdown for good text/code display
+    interactive=False,
+    lines=20
+)
+# --- END OF REPLACEMENT ---
 
 
     # Link components - ensure outputs match the function's yield tuple order
     submit_button.click(
         fn=run_agent_wrapper,
         inputs=query_input,
-        outputs=[reasoning_output_chatbot, final_answer_output] # Output to Chatbot and Textbox
+        outputs=[reasoning_output, final_answer_output]
     )
     # --- END OF CHANGE ---