Update app.py

app.py

@@ -10,6 +10,9 @@ import contextlib
 import sys
 import traceback
 import os
+import threading # <<< ADDED
+import queue     # <<< ADDED
+import time      # <<< ADDED
 
 @tool
 def parse_height_from_text(
@@ -229,94 +232,171 @@ else:
     )
     print(initialization_error_message)
     # height_agent is already None
+    
+# --- ADD THIS HELPER CLASS ---
+class WritableQueue:
+    """A file-like object that writes messages to a queue."""
+    def __init__(self, q):
+        self.queue = q
+
+    def write(self, message):
+        # Only put non-empty messages on the queue
+        if message.strip():
+            self.queue.put(message)
+
+    def flush(self):
+        # Required for file-like objects, but does nothing here
+        pass
+# --- END OF HELPER CLASS ---
+
+# --- REPLACE THE EXISTING run_agent_wrapper FUNCTION WITH THIS ---
+
+def agent_thread_func(agent, query, log_queue, result_queue):
+    """Function to run the agent in a separate thread and capture output."""
+    try:
+        # Create a WritableQueue instance for stdout redirection
+        stdout_writer = WritableQueue(log_queue)
+        # Redirect stdout within this thread
+        with contextlib.redirect_stdout(stdout_writer):
+            # Run the agent (prints will go to stdout_writer -> log_queue)
+            final_result = agent.run(query)
+            result_queue.put(final_result) # Put the final result in the result queue
+    except Exception as e:
+        # If an error occurs, print it to the log and put it in the result queue
+        tb_str = traceback.format_exc()
+        print(f"\n--- ERROR IN AGENT THREAD ---\n{e}\n{tb_str}")
+        result_queue.put(e) # Put the exception object itself
+    finally:
+        # Signal that logging is finished by putting None in the log queue
+        log_queue.put(None)
 
-# --- Wrapper Function to Run Agent and Capture Output ---
-def run_agent_wrapper(query: str) -> Tuple[str, str]:
+# Make the main Gradio function a generator
+def run_agent_wrapper(query: str) -> Iterator[Tuple[List[Tuple[str, str]], str]]:
     """
-    Runs the height_agent and captures its stdout (reasoning steps).
-    Returns (reasoning_log, final_answer).
+    Runs the agent in a thread, captures stdout via a queue, and yields updates
+    for Gradio streaming. Uses Chatbot format for reasoning.
+    Returns Iterator yielding: (chatbot_history, final_answer_status)
     """
-    # Access the global variables
-    global height_agent, initialization_error_message
-
     if height_agent is None:
-        # If agent initialization failed, return the stored error message
-        return (initialization_error_message or "Agent not initialized (unknown error).",
-                "Agent failed to initialize. See reasoning log for details.")
+        error_msg = initialization_error_message or "Agent not initialized."
+        yield ([(None, error_msg)], "Error: Agent not initialized.")
+        return # Stop the generator
+
+    log_queue = queue.Queue()
+    result_queue = queue.Queue()
+    chatbot_history = [] # Start with empty 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 (chatbot_history, final_answer)
+
+    # Start the agent thread
+    thread = threading.Thread(
+        target=agent_thread_func,
+        args=(height_agent, query, log_queue, result_queue)
+    )
+    thread.start()
 
-    print(f"\n--- Running agent for query: '{query}' ---") # Log to console
-    log_stream = io.StringIO()
-    final_answer = "Agent execution did not complete." # Default message
+    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
+            current_log_message += log_line
+            # Update the chatbot history: Replace the last message or add a new one
+            # Simple approach: always update a single entry representing the log
+            if chatbot_history and chatbot_history[-1][0] is None: # Check if last entry is from "Bot" (None for user)
+                 chatbot_history[-1] = (None, current_log_message) # Update last bot message
+            else:
+                 chatbot_history.append((None, current_log_message)) # Add new bot message if history is empty or last was user
+
+            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
+                while not log_queue.empty():
+                     log_line = log_queue.get_nowait()
+                     if log_line: current_log_message += log_line
+                if chatbot_history and chatbot_history[-1][0] is None:
+                     chatbot_history[-1] = (None, current_log_message + "\nError: Agent stopped unexpectedly.")
+                else:
+                     chatbot_history.append((None, current_log_message + "\nError: Agent stopped unexpectedly."))
+                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--- EXECUTION ERROR ---\n{final_result}"
+        current_log_message += error_log
+        if chatbot_history and chatbot_history[-1][0] is None:
+             chatbot_history[-1] = (None, current_log_message)
+        else:
+             chatbot_history.append((None, current_log_message))
+    else:
+        final_answer = final_result
 
-    try:
-        # Redirect stdout to capture prints from agent.run() (due to verbosity=3)
-        with contextlib.redirect_stdout(log_stream):
-            # Make sure to call the run method of the specific agent instance
-            final_answer = height_agent.run(query) # Pass the raw query
-            print("\n--- Agent execution finished successfully. ---") # Add marker to log
-    except Exception as e:
-        print(f"\n--- Error during agent execution wrapper: {e} ---") # Log to console
-        # Print exception details *into the captured log*
-        print("\n\n******** ERROR DURING EXECUTION ********\n", file=log_stream)
-        traceback.print_exc(file=log_stream)
-        final_answer = f"An error occurred during processing. See reasoning log. Error: {e}"
-    finally:
-        reasoning_log = log_stream.getvalue()
-        log_stream.close()
-        print("--- Finished capturing stdout. ---") # Log to console
+    # Final yield with the complete log and the final answer
+    yield (chatbot_history, final_answer)
 
-    return reasoning_log, final_answer
 # --- Build Gradio Interface Manually with gr.Blocks ---
 print("--- Building Gradio Interface with gr.Blocks ---")
 
 # Make sure theme is applied correctly if desired
 # theme = gr.themes.Default() # Or another theme
 # with gr.Blocks(theme=theme, css="footer {visibility: hidden}") as demo:
-with gr.Blocks(css="footer {visibility: hidden}") as demo: # Hides the default footer
+# --- MODIFY THE gr.Blocks SECTION ---
+
+with gr.Blocks(css="footer {visibility: hidden}") as demo:
     gr.Markdown("# Height Comparison Agent")
     gr.Markdown("Enter your height (e.g., '180 cm', '5ft 11in') to find characters/figures of similar height.")
 
     with gr.Row():
         with gr.Column(scale=1):
-            query_input = gr.Textbox(
-                label="Your Query (including height)",
-                placeholder="e.g., I am 175cm tall",
-                lines=2 # Allow slightly more room for input
-            )
+            query_input = gr.Textbox(label="Your Query (including height)", placeholder="e.g., I am 175cm tall")
             submit_button = gr.Button("Compare Heights", variant="primary")
         with gr.Column(scale=2):
-            final_answer_output = gr.Textbox(
-                label="Final Answer",
-                interactive=False,
-                lines=5
-            )
+            # Keep the Textbox for the final answer separate
+            final_answer_output = gr.Textbox(label="Final Answer", interactive=False, lines=5)
 
     gr.Markdown("## Agent Reasoning Steps")
-    # Use gr.Code for better formatting of logs, especially if they contain code blocks
-    reasoning_output = gr.Code(
+    # --- CHANGE THIS ---
+    # reasoning_output = gr.Code(label="Reasoning Log", language="markdown", interactive=False, lines=20)
+    reasoning_output_chatbot = gr.Chatbot(
         label="Reasoning Log",
-        language="markdown", # Use 'markdown' if logs might contain markdown
-        interactive=False,
-        lines=20
+        height=500 # Set a height to enable scrolling
     )
+    # --- END OF CHANGE ---
 
-    # Link components: When button is clicked, call wrapper, update outputs
+
+    # --- CHANGE THIS ---
+    # Link components - ensure outputs match the function's yield tuple order
     submit_button.click(
-        fn=run_agent_wrapper,          # Function to call
-        inputs=[query_input],          # Component(s) providing input
-        outputs=[reasoning_output, final_answer_output] # Components to update
-        # Ensure the order matches the return tuple from run_agent_wrapper: (log, answer)
+        fn=run_agent_wrapper,
+        inputs=query_input,
+        outputs=[reasoning_output_chatbot, final_answer_output] # Output to Chatbot and Textbox
     )
+    # --- END OF CHANGE ---
+
+# --- Launch Gradio (no change needed here) ---
+print("--- Starting Gradio Interface ---")
+demo.launch()
 
-    # Add an example input
-    gr.Examples(
-        examples=[
-            "I am 188cm tall",
-            "How tall is someone who is 5 foot 8 inches?",
-            "My height is 1.65m",
-        ],
-        inputs=query_input
-    )
-    # --- Launch Gradio ---
-print("--- Launching Gradio demo ---")
-demo.launch() # ssr=False recommended, share=True not needed for Spaces