diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index 7e0fe78..a7dfce5 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -24,6 +24,7 @@ from ra_aid.prompts import ( IMPLEMENTATION_PROMPT, SUMMARY_PROMPT ) +from ra_aid.exceptions import TaskCompletedException def parse_arguments(): parser = argparse.ArgumentParser( @@ -237,78 +238,85 @@ def validate_environment(): def main(): """Main entry point for the ra-aid command line tool.""" - validate_environment() - args = parse_arguments() - base_task = args.message - config = { - "configurable": { - "thread_id": "abc123" - }, - "recursion_limit": 100, - "research_only": args.research_only, - "cowboy_mode": args.cowboy_mode - } - - # Store config in global memory for access by is_informational_query - _global_memory['config'] = config + try: + validate_environment() + args = parse_arguments() + base_task = args.message + config = { + "configurable": { + "thread_id": "abc123" + }, + "recursion_limit": 100, + "research_only": args.research_only, + "cowboy_mode": args.cowboy_mode + } + + # Store config in global memory for access by is_informational_query + _global_memory['config'] = config - # Run research stage - print_stage_header("RESEARCH STAGE") - research_prompt = f"""User query: {base_task} --keep it simple + # Run research stage + print_stage_header("RESEARCH STAGE") + research_prompt = f"""User query: {base_task} --keep it simple {RESEARCH_PROMPT} Be very thorough in your research and emit lots of snippets, key facts. If you take more than a few steps, be eager to emit research subtasks.{'' if args.research_only else ' Only request implementation if the user explicitly asked for changes to be made.'}""" - while True: try: - for chunk in research_agent.stream( - {"messages": [HumanMessage(content=research_prompt)]}, - config - ): - print_agent_output(chunk) - break - except ChatAnthropic.InternalServerError as e: - print(f"Encountered Anthropic Internal Server Error: {e}. Retrying...") - continue + while True: + try: + for chunk in research_agent.stream( + {"messages": [HumanMessage(content=research_prompt)]}, + config + ): + print_agent_output(chunk) + break + except ChatAnthropic.InternalServerError as e: + print(f"Encountered Anthropic Internal Server Error: {e}. Retrying...") + continue + except TaskCompletedException as e: + print_stage_header("TASK COMPLETED") + raise # Re-raise to be caught by outer handler - # Run any research subtasks - run_research_subtasks(base_task, config) - - # For informational queries, summarize findings - if is_informational_query(): - summarize_research_findings(base_task, config) - else: - # Only proceed with planning and implementation if not an informational query - print_stage_header("PLANNING STAGE") - planning_prompt = PLANNING_PROMPT.format( - research_notes=get_memory_value('research_notes'), - key_facts=get_memory_value('key_facts'), - key_snippets=get_memory_value('key_snippets'), - base_task=base_task, - related_files="\n".join(related_files) - ) + # Run any research subtasks + run_research_subtasks(base_task, config) + + # For informational queries, summarize findings + if is_informational_query(): + summarize_research_findings(base_task, config) + else: + # Only proceed with planning and implementation if not an informational query + print_stage_header("PLANNING STAGE") + planning_prompt = PLANNING_PROMPT.format( + research_notes=get_memory_value('research_notes'), + key_facts=get_memory_value('key_facts'), + key_snippets=get_memory_value('key_snippets'), + base_task=base_task, + related_files="\n".join(related_files) + ) - # Run planning agent - while True: - try: - for chunk in planning_agent.stream( - {"messages": [HumanMessage(content=planning_prompt)]}, - config - ): - print_agent_output(chunk) - break - except ChatAnthropic.InternalServerError as e: - print(f"Encountered Anthropic Internal Server Error: {e}. Retrying...") - continue + # Run planning agent + while True: + try: + for chunk in planning_agent.stream( + {"messages": [HumanMessage(content=planning_prompt)]}, + config + ): + print_agent_output(chunk) + break + except ChatAnthropic.InternalServerError as e: + print(f"Encountered Anthropic Internal Server Error: {e}. Retrying...") + continue - # Run implementation stage with task-specific agents - run_implementation_stage( - base_task, - get_memory_value('tasks'), - get_memory_value('plan'), - related_files - ) + # Run implementation stage with task-specific agents + run_implementation_stage( + base_task, + get_memory_value('tasks'), + get_memory_value('plan'), + related_files + ) + except TaskCompletedException: + sys.exit(0) if __name__ == "__main__": main() diff --git a/ra_aid/exceptions.py b/ra_aid/exceptions.py new file mode 100644 index 0000000..1c4d8b2 --- /dev/null +++ b/ra_aid/exceptions.py @@ -0,0 +1,3 @@ +class TaskCompletedException(Exception): + """Raised when a one-shot task has been completed.""" + pass diff --git a/ra_aid/prompts.py b/ra_aid/prompts.py index cd1285c..6ebedfe 100644 --- a/ra_aid/prompts.py +++ b/ra_aid/prompts.py @@ -61,6 +61,22 @@ No Planning or Problem-Solving Do not provide advice or commentary on the project’s future. You must remain strictly within the bounds of describing what currently exists. + +Single-Shot Task Detection + + Autonomously determine if a task can be completed immediately without further planning: + - Simple informational queries that can be answered directly from research + - Requests that don't require complex analysis or implementation + - Cases where further planning would not add value + - Situations where immediate response meets all user requirements + + If you determine a task can be completed in a single shot: + 1. Complete the necessary research to fully answer the query + 2. Document your findings using emit_research_notes + 3. Call one_shot_completed() to immediately conclude the task + + Only use single-shot completion when you are confident no implementation or further planning is needed. + Thoroughness and Completeness If this is determined to be a new/empty project (no code or files), state that and stop. diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index ed0ad3e..0ab41aa 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -1,4 +1,5 @@ from typing import Dict, List, Any, Union, TypedDict, Optional, Sequence +from ra_aid.exceptions import TaskCompletedException from rich.console import Console from rich.markdown import Markdown from rich.panel import Panel @@ -238,6 +239,21 @@ def delete_key_snippets(snippet_ids: List[int]) -> str: return "Snippets deleted." +@tool("one_shot_completed") +def one_shot_completed(message: str) -> str: + """Signal that a one-shot task has been completed and execution should stop. + + Args: + message: Completion message to display + + Raises: + TaskCompletedException: Always raised to stop execution + + Returns: + Never returns, always raises exception + """ + raise TaskCompletedException(message) + def get_memory_value(key: str) -> str: """Get a value from global memory.