Handle keyboard interrupt (ctrl-c.)
This commit is contained in:
AI Christianson
parent
5ad8aeb013
commit
c4efac339a
|
|
@ -247,104 +247,109 @@ def run_research_subtasks(base_task: str, config: dict, model, expert_enabled: b
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
"""Main entry point for the ra-aid command line tool."""
|
"""Main entry point for the ra-aid command line tool."""
|
||||||
args = parse_arguments()
|
try:
|
||||||
expert_enabled, expert_missing = validate_environment(args) # Will exit if main env vars missing
|
args = parse_arguments()
|
||||||
|
expert_enabled, expert_missing = validate_environment(args) # Will exit if main env vars missing
|
||||||
if expert_missing:
|
|
||||||
console.print(Panel(
|
if expert_missing:
|
||||||
f"[yellow]Expert tools disabled due to missing configuration:[/yellow]\n" +
|
console.print(Panel(
|
||||||
"\n".join(f"- {m}" for m in expert_missing) +
|
f"[yellow]Expert tools disabled due to missing configuration:[/yellow]\n" +
|
||||||
"\nSet the required environment variables or args to enable expert mode.",
|
"\n".join(f"- {m}" for m in expert_missing) +
|
||||||
title="Expert Tools Disabled",
|
"\nSet the required environment variables or args to enable expert mode.",
|
||||||
style="yellow"
|
title="Expert Tools Disabled",
|
||||||
))
|
style="yellow"
|
||||||
|
))
|
||||||
# Create the base model after validation
|
|
||||||
model = initialize_llm(args.provider, args.model)
|
# Create the base model after validation
|
||||||
|
model = initialize_llm(args.provider, args.model)
|
||||||
|
|
||||||
# Validate message is provided
|
# Validate message is provided
|
||||||
if not args.message:
|
if not args.message:
|
||||||
print_error("--message is required")
|
print_error("--message is required")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
# Store expert provider and model in config
|
||||||
|
_global_memory['config']['expert_provider'] = args.expert_provider
|
||||||
|
_global_memory['config']['expert_model'] = args.expert_model
|
||||||
|
|
||||||
base_task = args.message
|
# Run research stage
|
||||||
config = {
|
print_stage_header("Research Stage")
|
||||||
"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
|
|
||||||
|
|
||||||
# Store expert provider and model in config
|
|
||||||
_global_memory['config']['expert_provider'] = args.expert_provider
|
|
||||||
_global_memory['config']['expert_model'] = args.expert_model
|
|
||||||
|
|
||||||
# Run research stage
|
|
||||||
print_stage_header("Research Stage")
|
|
||||||
|
|
||||||
# Create research agent
|
|
||||||
research_agent = create_react_agent(
|
|
||||||
model,
|
|
||||||
get_research_tools(
|
|
||||||
research_only=_global_memory.get('config', {}).get('research_only', False),
|
|
||||||
expert_enabled=expert_enabled,
|
|
||||||
human_interaction=args.hil
|
|
||||||
),
|
|
||||||
checkpointer=research_memory
|
|
||||||
)
|
|
||||||
|
|
||||||
expert_section = EXPERT_PROMPT_SECTION_RESEARCH if expert_enabled else ""
|
|
||||||
human_section = HUMAN_PROMPT_SECTION_RESEARCH if args.hil else ""
|
|
||||||
research_prompt = RESEARCH_PROMPT.format(
|
|
||||||
expert_section=expert_section,
|
|
||||||
human_section=human_section,
|
|
||||||
base_task=base_task,
|
|
||||||
research_only_note='' if args.research_only else ' Only request implementation if the user explicitly asked for changes to be made.'
|
|
||||||
)
|
|
||||||
|
|
||||||
# Run research agent
|
|
||||||
run_agent_with_retry(research_agent, research_prompt, config)
|
|
||||||
|
|
||||||
# Run any research subtasks
|
|
||||||
run_research_subtasks(base_task, config, model, expert_enabled=expert_enabled)
|
|
||||||
|
|
||||||
# Proceed with planning and implementation if not an informational query
|
|
||||||
if not is_informational_query():
|
|
||||||
print_stage_header("Planning Stage")
|
|
||||||
|
|
||||||
# Create planning agent
|
# Create research agent
|
||||||
planning_agent = create_react_agent(model, get_planning_tools(expert_enabled=expert_enabled), checkpointer=planning_memory)
|
research_agent = create_react_agent(
|
||||||
|
model,
|
||||||
expert_section = EXPERT_PROMPT_SECTION_PLANNING if expert_enabled else ""
|
get_research_tools(
|
||||||
human_section = HUMAN_PROMPT_SECTION_PLANNING if args.hil else ""
|
research_only=_global_memory.get('config', {}).get('research_only', False),
|
||||||
planning_prompt = PLANNING_PROMPT.format(
|
expert_enabled=expert_enabled,
|
||||||
|
human_interaction=args.hil
|
||||||
|
),
|
||||||
|
checkpointer=research_memory
|
||||||
|
)
|
||||||
|
|
||||||
|
expert_section = EXPERT_PROMPT_SECTION_RESEARCH if expert_enabled else ""
|
||||||
|
human_section = HUMAN_PROMPT_SECTION_RESEARCH if args.hil else ""
|
||||||
|
research_prompt = RESEARCH_PROMPT.format(
|
||||||
expert_section=expert_section,
|
expert_section=expert_section,
|
||||||
human_section=human_section,
|
human_section=human_section,
|
||||||
base_task=base_task,
|
base_task=base_task,
|
||||||
research_notes=get_memory_value('research_notes'),
|
|
||||||
related_files="\n".join(get_related_files()),
|
|
||||||
key_facts=get_memory_value('key_facts'),
|
|
||||||
key_snippets=get_memory_value('key_snippets'),
|
|
||||||
research_only_note='' if args.research_only else ' Only request implementation if the user explicitly asked for changes to be made.'
|
research_only_note='' if args.research_only else ' Only request implementation if the user explicitly asked for changes to be made.'
|
||||||
)
|
)
|
||||||
|
|
||||||
# Run planning agent
|
# Run research agent
|
||||||
run_agent_with_retry(planning_agent, planning_prompt, config)
|
run_agent_with_retry(research_agent, research_prompt, config)
|
||||||
|
|
||||||
|
# Run any research subtasks
|
||||||
|
run_research_subtasks(base_task, config, model, expert_enabled=expert_enabled)
|
||||||
|
|
||||||
|
# Proceed with planning and implementation if not an informational query
|
||||||
|
if not is_informational_query():
|
||||||
|
print_stage_header("Planning Stage")
|
||||||
|
|
||||||
|
# Create planning agent
|
||||||
|
planning_agent = create_react_agent(model, get_planning_tools(expert_enabled=expert_enabled), checkpointer=planning_memory)
|
||||||
|
|
||||||
|
expert_section = EXPERT_PROMPT_SECTION_PLANNING if expert_enabled else ""
|
||||||
|
human_section = HUMAN_PROMPT_SECTION_PLANNING if args.hil else ""
|
||||||
|
planning_prompt = PLANNING_PROMPT.format(
|
||||||
|
expert_section=expert_section,
|
||||||
|
human_section=human_section,
|
||||||
|
base_task=base_task,
|
||||||
|
research_notes=get_memory_value('research_notes'),
|
||||||
|
related_files="\n".join(get_related_files()),
|
||||||
|
key_facts=get_memory_value('key_facts'),
|
||||||
|
key_snippets=get_memory_value('key_snippets'),
|
||||||
|
research_only_note='' if args.research_only else ' Only request implementation if the user explicitly asked for changes to be made.'
|
||||||
|
)
|
||||||
|
|
||||||
# Run implementation stage with task-specific agents
|
# Run planning agent
|
||||||
run_implementation_stage(
|
run_agent_with_retry(planning_agent, planning_prompt, config)
|
||||||
base_task,
|
|
||||||
get_memory_value('tasks'),
|
# Run implementation stage with task-specific agents
|
||||||
get_memory_value('plan'),
|
run_implementation_stage(
|
||||||
get_related_files(),
|
base_task,
|
||||||
model,
|
get_memory_value('tasks'),
|
||||||
expert_enabled=expert_enabled
|
get_memory_value('plan'),
|
||||||
)
|
get_related_files(),
|
||||||
|
model,
|
||||||
|
expert_enabled=expert_enabled
|
||||||
|
)
|
||||||
|
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
console.print("\n[red]Operation cancelled by user[/red]")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue