Update app.py

app.py

@@ -5,7 +5,6 @@ import gradio as gr
 from multiprocessing import freeze_support
 import importlib
 import inspect
-import json
 
 # === Fix path to include src/txagent
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
@@ -37,6 +36,14 @@ question_examples = [
     ["What treatment options exist for HER2+ breast cancer resistant to trastuzumab?"]
 ]
 
+# === Helper: Add collapsible formatting
+def format_collapsible_response(content):
+    return (
+        f"<details style='border: 1px solid #ccc; padding: 8px; margin-top: 8px;'>"
+        f"<summary style='font-weight: bold;'>Answer</summary>"
+        f"<div style='margin-top: 8px;'>{content}</div></details>"
+    )
+
 # === UI creation
 def create_ui(agent):
     with gr.Blocks() as demo:
@@ -69,26 +76,12 @@ def create_ui(agent):
 
             for update in generator:
                 formatted_messages = []
-
                 for m in update:
-                    role = m.get("role")
-                    content = m.get("content", "")
-
-                    # 🧠 Pretty-print tool calls
-                    if "[TOOL_CALLS]" in content:
-                        try:
-                            thought, tools = content.split("[TOOL_CALLS]", 1)
-                            tools_json = json.loads(tools.strip())
-                            pretty_tools = json.dumps(tools_json, indent=2)
-
-                            # Style it like a ChatGPT output
-                            content = f"**🧠 Assistant reasoning:**\n\n{thought.strip()}\n\n" \
-                                      f"**🔧 Tool Calls:**\n\n```json\n{pretty_tools}\n```"
-                        except Exception as e:
-                            content = f"{content}\n\n⚠️ Failed to format tool calls: {e}"
-
+                    role = m["role"] if isinstance(m, dict) else getattr(m, "role", "assistant")
+                    content = m["content"] if isinstance(m, dict) else getattr(m, "content", "")
+                    if role == "assistant":
+                        content = format_collapsible_response(content)
                     formatted_messages.append({"role": role, "content": content})
-
                 yield formatted_messages
 
         # === Trigger handlers
@@ -123,7 +116,7 @@ if __name__ == "__main__":
             enable_checker=True,
             step_rag_num=10,
             seed=100,
-            additional_default_tools=["DirectResponse", "RequireClarification"]
+            additional_default_tools=[]  # Removed DirectResponse/RequireClarification to avoid errors
         )
         agent.init_model()