update

app.py

@@ -36,6 +36,14 @@ DATASET_DIR.mkdir(parents=True, exist_ok=True)
 global QUERY_ID
 QUERY_ID = None
 
+TOOL_NAME_MAPPING = {
+    "Generalist_Solution_Generator_Tool": "Base_Generator_Tool",
+    "Ground_Google_Search_Tool": "Google_Search_Tool",
+    "Python_Code_Generator_Tool": "Python_Coder_Tool",
+    "Web_RAG_Search_Tool": "Web_Search_Tool",
+    "Wikipedia_RAG_Search_Tool": "Wikipedia_Search_Tool"
+}
+
 # Enable scheduler to record data to HuggingFace dataset
 # scheduler = None
 scheduler = CommitScheduler(
@@ -266,13 +274,20 @@ class Solver:
 
         # [Step 4] Query Analysis
         query_analysis = self.planner.analyze_query(user_query, img_path)
-        json_data["query_analysis"] = query_analysis
-        query_analysis = query_analysis.replace("Concise Summary:", "**Concise Summary:**\n")
-        query_analysis = query_analysis.replace("Required Skills:", "**Required Skills:**")
-        query_analysis = query_analysis.replace("Relevant Tools:", "**Relevant Tools:**")
-        query_analysis = query_analysis.replace("Additional Considerations:", "**Additional Considerations:**")
+        json_data["query_analysis"] = query_analysis # TODO: update
+
+        # Format the query analysis for display
+        query_analysis_display = query_analysis.replace("Concise Summary:", "**Concise Summary:**\n")
+        query_analysis_display = query_analysis_display.replace("Required Skills:", "**Required Skills:**")
+        query_analysis_display = query_analysis_display.replace("Relevant Tools:", "**Relevant Tools:**")
+        query_analysis_display = query_analysis_display.replace("Additional Considerations:", "**Additional Considerations:**")
+
+        # Map tool names in query analysis for display
+        for original_name, display_name in TOOL_NAME_MAPPING.items():
+            query_analysis_display = query_analysis_display.replace(original_name, display_name)
+
         messages.append(ChatMessage(role="assistant",
-                                    content=f"{query_analysis}",
+                                    content=f"{query_analysis_display}",
                                     metadata={"title": "### 🔎 Step 0: Query Analysis"}))
         yield messages
 
@@ -297,7 +312,7 @@ class Solver:
             next_step = self.planner.generate_next_step(
                 user_query, img_path, query_analysis, self.memory, step_count, self.max_steps, json_data
             )
-            context, sub_goal, tool_name = self.planner.extract_context_subgoal_and_tool(next_step)
+            context, sub_goal, tool_name = self.planner.extract_context_subgoal_and_tool(next_step) # TODO: update
             step_data = {
                 "step_count": step_count,
                 "context": context,
@@ -308,17 +323,27 @@ class Solver:
             save_module_data(QUERY_ID, f"step_{step_count}_action_prediction", step_data)
 
             # Display the step information
+            display_tool_name = TOOL_NAME_MAPPING.get(tool_name, tool_name)
+
+            # Map tool names in context and sub_goal for display
+            context_display = context if context else ""
+            sub_goal_display = sub_goal if sub_goal else ""
+            for original_name, display_name in TOOL_NAME_MAPPING.items():
+                context_display = context_display.replace(original_name, display_name)
+                sub_goal_display = sub_goal_display.replace(original_name, display_name)
+
             messages.append(ChatMessage(
                 role="assistant",
-                content=f"**Context:** {context}\n\n**Sub-goal:** {sub_goal}\n\n**Tool:** `{tool_name}`",
-                metadata={"title": f"### 🎯 Step {step_count}: Action Prediction ({tool_name})"}))
+                content=f"**Context:** {context_display}\n\n**Sub-goal:** {sub_goal_display}\n\n**Tool:** `{display_tool_name}`",
+                metadata={"title": f"### 🎯 Step {step_count}: Action Prediction ({display_tool_name})"}))
             yield messages
 
             # Handle tool execution or errors
             if tool_name not in self.planner.available_tools:
+                display_tool_name = TOOL_NAME_MAPPING.get(tool_name, tool_name)
                 messages.append(ChatMessage(
-                    role="assistant", 
-                    content=f"⚠️ Error: Tool '{tool_name}' is not available."))
+                    role="assistant",
+                    content=f"⚠️ Error: Tool '{display_tool_name}' is not available."))
                 yield messages
                 continue
 
@@ -331,10 +356,11 @@ class Solver:
             result = make_json_serializable_truncated(result)
 
             # Display the ommand generation information
+            display_tool_name = TOOL_NAME_MAPPING.get(tool_name, tool_name)
             messages.append(ChatMessage(
                 role="assistant",
                 content=f"**Command:**\n```python\n{command}\n```",
-                metadata={"title": f"### 📋 Step {step_count}: Command Generation ({tool_name})"}))
+                metadata={"title": f"### 📋 Step {step_count}: Command Generation ({display_tool_name})"}))
             yield messages
 
             # Save the command generation data
@@ -347,11 +373,18 @@ class Solver:
             save_module_data(QUERY_ID, f"step_{step_count}_command_generation", command_generation_data)
             
             # Display the command execution result
+            display_tool_name = TOOL_NAME_MAPPING.get(tool_name, tool_name)
+
+            # Map tool names in result for display
+            result_json_str = json.dumps(result, indent=4)
+            for original_name, display_name in TOOL_NAME_MAPPING.items():
+                result_json_str = result_json_str.replace(original_name, display_name)
+
             messages.append(ChatMessage(
                 role="assistant",
-                content=f"**Result:**\n```json\n{json.dumps(result, indent=4)}\n```",
+                content=f"**Result:**\n```json\n{result_json_str}\n```",
                 # content=f"**Result:**\n```json\n{result}\n```",
-                metadata={"title": f"### ⚡ Step {step_count}: Command Execution ({tool_name})"}))
+                metadata={"title": f"### ⚡ Step {step_count}: Command Execution ({display_tool_name})"}))
             yield messages
 
             # Save the command execution data
@@ -362,7 +395,7 @@ class Solver:
             save_module_data(QUERY_ID, f"step_{step_count}_command_execution", command_execution_data)
 
             # [Step 8] Memory update and stopping condition
-            self.memory.add_action(step_count, tool_name, sub_goal, command, result)
+            self.memory.add_action(step_count, tool_name, sub_goal, command, result) # TODO: do not update here
             stop_verification = self.planner.verificate_context(user_query, img_path, query_analysis, self.memory, step_count, json_data)
             context_verification, conclusion = self.planner.extract_conclusion(stop_verification)
 
@@ -372,13 +405,18 @@ class Solver:
                 "conclusion": conclusion,
                 "time": round(time.time() - start_time, 5)
             }
-            save_module_data(QUERY_ID, f"step_{step_count}_context_verification", context_verification_data)    
+            save_module_data(QUERY_ID, f"step_{step_count}_context_verification", context_verification_data)
+
+            # Display the context verification result # TODO: update context_verification
+            # Map tool names in context verification for display
+            context_verification_display = context_verification if context_verification else ""
+            for original_name, display_name in TOOL_NAME_MAPPING.items():
+                context_verification_display = context_verification_display.replace(original_name, display_name)
 
-            # Display the context verification result
             conclusion_emoji = "✅" if conclusion == 'STOP' else "🛑"
             messages.append(ChatMessage(
                 role="assistant",
-                content=f"**Analysis:**\n{context_verification}\n\n**Conclusion:** `{conclusion}` {conclusion_emoji}",
+                content=f"**Analysis:**\n{context_verification_display}\n\n**Conclusion:** `{conclusion}` {conclusion_emoji}",
                 metadata={"title": f"### 🤖 Step {step_count}: Context Verification"}))
             yield messages
 
@@ -388,8 +426,14 @@ class Solver:
         # Step 7: Generate Final Output (if needed)
         if 'direct' in self.output_types:
             messages.append(ChatMessage(role="assistant", content="<br>"))
-            direct_output = self.planner.generate_direct_output(user_query, img_path, self.memory)
-            messages.append(ChatMessage(role="assistant", content=f"### 🎉 Final Answer:\n{direct_output}"))
+            direct_output = self.planner.generate_direct_output(user_query, img_path, self.memory) # TODO: update
+
+            # Map tool names in direct output for display
+            direct_output_display = direct_output if direct_output else ""
+            for original_name, display_name in TOOL_NAME_MAPPING.items():
+                direct_output_display = direct_output_display.replace(original_name, display_name)
+
+            messages.append(ChatMessage(role="assistant", content=f"### 🎉 Final Answer:\n{direct_output_display}"))
             yield messages
 
             # Save the direct output data
@@ -545,7 +589,9 @@ def main(args):
         [Paper](https://arxiv.org/abs/2510.05592) |
         [GitHub](https://github.com/lupantech/AgentFlow) |
 
-        > ⏳ **Note:** The first query may take ~20 seconds to initialize AgentFlow. Subsequent queries will be supper fast.
+        > ⏳ **Note:** The first query may take ~20 seconds to initialize AgentFlow (due to heavy traffic). Subsequent queries will be super fast.
+        >
+        > 💡 **Tip:** If the wait time is too long, please try again later.
         """)
 
         with gr.Row():