diff --git a/ra_aid/agent_utils.py b/ra_aid/agent_utils.py index 6be9f99..11ed096 100644 --- a/ra_aid/agent_utils.py +++ b/ra_aid/agent_utils.py @@ -218,16 +218,6 @@ def run_agent_with_retry(agent, prompt: str, config: dict) -> Optional[str]: config ): print_agent_output(chunk) - - # Check for task completion after each chunk - if _global_memory.get('task_completed'): - completion_msg = _global_memory.get('completion_message', 'Task was completed successfully.') - console.print(Panel( - Markdown(completion_msg), - title="✅ Task Completed", - style="green" - )) - return completion_msg break except (InternalServerError, APITimeoutError, RateLimitError, APIError) as e: if attempt == max_retries - 1: diff --git a/ra_aid/tool_configs.py b/ra_aid/tool_configs.py index 17c4c9f..6ad1e61 100644 --- a/ra_aid/tool_configs.py +++ b/ra_aid/tool_configs.py @@ -6,7 +6,8 @@ from ra_aid.tools import ( emit_key_snippets, delete_key_snippets, delete_tasks, request_implementation, read_file_tool, fuzzy_find_project_files, ripgrep_search, list_directory_tree, - swap_task_order, monorepo_detected, existing_project_detected, ui_detected + swap_task_order, monorepo_detected, existing_project_detected, ui_detected, + task_completed ) from ra_aid.tools.memory import one_shot_completed from ra_aid.tools.agent import request_research, request_task_implementation @@ -94,6 +95,9 @@ def get_implementation_tools(expert_enabled: bool = True) -> list: # Add modification tools since it's not research-only tools.extend(MODIFICATION_TOOLS) + tools.extend([ + task_completed + ]) # Add expert tools if enabled if expert_enabled: diff --git a/ra_aid/tools/__init__.py b/ra_aid/tools/__init__.py index 7c6d17e..b6e2ecf 100644 --- a/ra_aid/tools/__init__.py +++ b/ra_aid/tools/__init__.py @@ -12,7 +12,7 @@ from .ripgrep import ripgrep_search from .memory import ( delete_tasks, emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts, request_implementation, delete_key_facts, - emit_key_snippets, delete_key_snippets, emit_related_files, swap_task_order + emit_key_snippets, delete_key_snippets, emit_related_files, swap_task_order, task_completed ) __all__ = [ @@ -41,5 +41,6 @@ __all__ = [ 'monorepo_detected', 'existing_project_detected', 'ui_detected', - 'ask_human' + 'ask_human', + 'task_completed' ] diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index 08acbee..48db433 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -299,6 +299,21 @@ def one_shot_completed(message: str) -> str: _global_memory['completion_message'] = message return message +@tool("task_completed") +def task_completed(message: str) -> str: + """Mark the current task as completed with a completion message. + + Args: + message: Message explaining how/why the task is complete + + Returns: + The completion message + """ + _global_memory['task_completed'] = True + _global_memory['completion_message'] = message + console.print(Panel(Markdown(message), title="✅ Task Completed")) + return "Completion noted." + def get_related_files() -> Set[str]: """Get the current set of related files.