Update app.py

app.py

@@ -249,6 +249,7 @@ class WritableQueue:
         pass
 # --- END OF HELPER CLASS ---
 
+# --- REPLACE THE EXISTING run_agent_wrapper FUNCTION WITH THIS ---
 # --- REPLACE THE EXISTING run_agent_wrapper FUNCTION WITH THIS ---
 
 def agent_thread_func(agent, query, log_queue, result_queue):
@@ -270,26 +271,31 @@ def agent_thread_func(agent, query, log_queue, result_queue):
         # Signal that logging is finished by putting None in the log queue
         log_queue.put(None)
 
-# Make the main Gradio function a generator
-def run_agent_wrapper(query: str) -> Iterator[Tuple[List[Tuple[str, str]], str]]:
+# Generator function for Gradio streaming
+# REMOVED the return type hint -> Iterator[...]
+def run_agent_wrapper(query: str):
     """
     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)
+    for Gradio streaming. Uses Chatbot 'messages' format for reasoning.
+    Yields: (chatbot_history, final_answer_status)
     """
     if height_agent is None:
         error_msg = initialization_error_message or "Agent not initialized."
-        yield ([(None, error_msg)], "Error: 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()
-    chatbot_history = [] # Start with empty history
+    # 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 (chatbot_history, final_answer)
+    # Yield empty history initially, or a starting message
+    yield ([], final_answer)
 
     # Start the agent thread
     thread = threading.Thread(
@@ -307,13 +313,12 @@ def run_agent_wrapper(query: str) -> Iterator[Tuple[List[Tuple[str, str]], str]]
                 break
 
             # Append new line to the current log message block
-            current_log_message += log_line
+            # 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
-            # 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
+            # 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
 
@@ -323,13 +328,15 @@ def run_agent_wrapper(query: str) -> Iterator[Tuple[List[Tuple[str, str]], str]]
             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."))
+                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
@@ -345,25 +352,25 @@ def run_agent_wrapper(query: str) -> Iterator[Tuple[List[Tuple[str, str]], str]]
     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}"
+        error_log = f"\n\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))
+        chatbot_history = [{"role": "assistant", "content": current_log_message}]
     else:
-        final_answer = final_result
+        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)
 
+# --- END OF REPLACEMENT ---
+
 # --- 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:
-# --- MODIFY THE gr.Blocks SECTION ---
 
 with gr.Blocks(css="footer {visibility: hidden}") as demo:
     gr.Markdown("# Height Comparison Agent")
@@ -381,13 +388,13 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
     # --- CHANGE THIS ---
     # reasoning_output = gr.Code(label="Reasoning Log", language="markdown", interactive=False, lines=20)
     reasoning_output_chatbot = gr.Chatbot(
-        label="Reasoning Log",
-        height=500 # Set a height to enable scrolling
+    label="Reasoning Log",
+    height=500,
+    type="messages" # <<< ADD THIS
+`)
     )
-    # --- END OF CHANGE ---
 
 
-    # --- CHANGE THIS ---
     # Link components - ensure outputs match the function's yield tuple order
     submit_button.click(
         fn=run_agent_wrapper,