Cleanup.

ra_aid/__main__.py

@@ -7,8 +7,8 @@ from langgraph.checkpoint.memory import MemorySaver
 from langgraph.prebuilt import create_react_agent
 from ra_aid.env import validate_environment
 from ra_aid.tools.memory import _global_memory, get_related_files, get_memory_value
-from ra_aid import print_stage_header, print_task_header, print_error
-from ra_aid.agent_utils import run_agent_with_retry, run_task_implementation_agent
+from ra_aid import print_stage_header, print_error
+from ra_aid.agent_utils import run_agent_with_retry
 from ra_aid.agent_utils import run_research_agent
 from ra_aid.prompts import (
     PLANNING_PROMPT,
@@ -20,7 +20,6 @@ from ra_aid.llm import initialize_llm
 
 from ra_aid.tool_configs import (
     get_planning_tools,
-    get_implementation_tools,
     get_chat_tools
 )
 

ra_aid/tools/memory.py

@@ -287,17 +287,14 @@ def one_shot_completed(message: str) -> str:
     
     Args:
         message: Completion message to display
-        
-    Returns:
-        Original message if task can be completed, or error message if there are
-        pending subtasks or implementation requests
     """
     if _global_memory.get('implementation_requested', False):
         return "Cannot complete in one shot - implementation was requested"
         
     _global_memory['task_completed'] = True
     _global_memory['completion_message'] = message
-    return message
+    console.print(Panel(Markdown(message), title="✅ One-Shot Task Completed"))
+    return "Completion noted."
 
 @tool("task_completed")
 def task_completed(message: str) -> str: