expert API key fallback

ra_aid/__main__.py

@@ -102,8 +102,8 @@ research_memory = MemorySaver()
 planning_memory = MemorySaver()
 implementation_memory = MemorySaver()
 
-def get_research_tools(research_only: bool = False, expert_enabled: bool = True, llm_enabled: bool = True) -> list:
+def get_research_tools(research_only: bool = False, expert_enabled: bool = True) -> list:
-    """Get the list of research tools based on mode and availability."""
+    """Get the list of research tools based on mode and whether expert is enabled."""
     tools = [
         list_directory_tree,
         emit_research_subtask,
@@ -119,16 +119,16 @@ def get_research_tools(research_only: bool = False, expert_enabled: bool = True,
         ripgrep_search
     ]
     
-    if expert_enabled and llm_enabled:
+    if expert_enabled:
         tools.append(emit_expert_context)
         tools.append(ask_expert)
     
-    if not research_only and llm_enabled:
+    if not research_only:
         tools.append(request_implementation)
     
     return tools
 
-def get_planning_tools(expert_enabled: bool = True, llm_enabled: bool = True) -> list:
+def get_planning_tools(expert_enabled: bool = True) -> list:
     tools = [
         list_directory_tree,
         emit_plan,
@@ -142,12 +142,12 @@ def get_planning_tools(expert_enabled: bool = True, llm_enabled: bool = True) ->
         fuzzy_find_project_files,
         ripgrep_search
     ]
-    if expert_enabled and llm_enabled:
+    if expert_enabled:
         tools.append(ask_expert)
         tools.append(emit_expert_context)
     return tools
 
-def get_implementation_tools(expert_enabled: bool = True, llm_enabled: bool = True) -> list:
+def get_implementation_tools(expert_enabled: bool = True) -> list:
     tools = [
         list_directory_tree,
         run_shell_command,
@@ -161,22 +161,25 @@ def get_implementation_tools(expert_enabled: bool = True, llm_enabled: bool = Tr
         fuzzy_find_project_files,
         ripgrep_search
     ]
-    if expert_enabled and llm_enabled:
+    if expert_enabled:
         tools.append(ask_expert)
         tools.append(emit_expert_context)
     return tools
 
 def is_informational_query() -> bool:
+    """Determine if the current query is informational based on implementation_requested state."""
     return _global_memory.get('config', {}).get('research_only', False) or not is_stage_requested('implementation')
 
 def is_stage_requested(stage: str) -> bool:
+    """Check if a stage has been requested to proceed."""
     if stage == 'implementation':
         return len(_global_memory.get('implementation_requested', [])) > 0
     return False
 
 def run_agent_with_retry(agent, prompt: str, config: dict):
+    """Run an agent with retry logic for internal server errors."""
     max_retries = 20
-    base_delay = 1
+    base_delay = 1  # Initial delay in seconds
     
     for attempt in range(max_retries):
         try:
@@ -190,19 +193,21 @@ def run_agent_with_retry(agent, prompt: str, config: dict):
             if attempt == max_retries - 1:
                 raise RuntimeError(f"Max retries ({max_retries}) exceeded. Last error: {str(e)}")
             
-            delay = base_delay * (2 ** attempt)
+            delay = base_delay * (2 ** attempt)  # Exponential backoff
             error_type = e.__class__.__name__
             print_error(f"Encountered {error_type}: {str(e)}. Retrying in {delay} seconds... (Attempt {attempt + 1}/{max_retries})")
             time.sleep(delay)
             continue
 
-def run_implementation_stage(base_task, tasks, plan, related_files, model, expert_enabled: bool, llm_enabled: bool):
+def run_implementation_stage(base_task, tasks, plan, related_files, model, expert_enabled: bool):
+    """Run implementation stage with a distinct agent for each task."""
     if not is_stage_requested('implementation'):
         print_stage_header("Implementation Stage Skipped")
         return
         
     print_stage_header("Implementation Stage")
     
+    # Get tasks directly from memory
     task_list = _global_memory['tasks']
     
     print_task_header(f"Found {len(task_list)} tasks to implement")
@@ -210,9 +215,13 @@ def run_implementation_stage(base_task, tasks, plan, related_files, model, exper
     for i, task in enumerate(task_list, 1):
         print_task_header(task)
         
+        # Create a unique memory instance for this task
         task_memory = MemorySaver()
-        task_agent = create_react_agent(model, get_implementation_tools(expert_enabled=expert_enabled, llm_enabled=llm_enabled), checkpointer=task_memory)
         
+        # Create a fresh agent for each task
+        task_agent = create_react_agent(model, get_implementation_tools(expert_enabled=expert_enabled), checkpointer=task_memory)
+        
+        # Construct task-specific prompt
         task_prompt = IMPLEMENTATION_PROMPT.format(
             plan=plan,
             key_facts=get_memory_value('key_facts'),
@@ -222,27 +231,29 @@ def run_implementation_stage(base_task, tasks, plan, related_files, model, exper
             base_task=base_task
         )
         
-        if llm_enabled:
+        # Run agent for this task
         run_agent_with_retry(task_agent, task_prompt, {"configurable": {"thread_id": "abc123"}, "recursion_limit": 100})
-        else:
-            console.print(Panel("[yellow]LLM is disabled, cannot implement tasks[/yellow]", title="No LLM Available"))
 
-def run_research_subtasks(base_task: str, config: dict, model, expert_enabled: bool, llm_enabled: bool):
+
+def run_research_subtasks(base_task: str, config: dict, model, expert_enabled: bool):
+    """Run research subtasks with separate agents."""
     subtasks = _global_memory.get('research_subtasks', [])
     if not subtasks:
         return
         
     print_stage_header("Research Subtasks")
     
+    # Get tools for subtask agents (excluding emit_research_subtask and implementation)
     research_only = _global_memory.get('config', {}).get('research_only', False)
     subtask_tools = [
-        t for t in get_research_tools(research_only=research_only, expert_enabled=expert_enabled, llm_enabled=llm_enabled)
+        t for t in get_research_tools(research_only=research_only, expert_enabled=expert_enabled)
         if t.name not in ['emit_research_subtask']
     ]
     
     for i, subtask in enumerate(subtasks, 1):
         print_task_header(f"Research Subtask {i}/{len(subtasks)}")
         
+        # Create fresh memory and agent for each subtask
         subtask_memory = MemorySaver()
         subtask_agent = create_react_agent(
             model,
@@ -250,19 +261,25 @@ def run_research_subtasks(base_task: str, config: dict, model, expert_enabled: b
             checkpointer=subtask_memory
         )
         
+        # Run the subtask agent
         subtask_prompt = f"Research Subtask: {subtask}\n\n{RESEARCH_PROMPT}"
-        
-        if llm_enabled:
         run_agent_with_retry(subtask_agent, subtask_prompt, config)
-        else:
+
-            console.print(Panel("[yellow]LLM is disabled, cannot perform LLM-based research[/yellow]", title="No LLM Available"))
 
 def validate_environment(args):
+    """Validate required environment variables and dependencies.
+    
+    Args:
+        args: The parsed command line arguments
+        
+    Returns:
+        (expert_enabled, missing_expert_info): Tuple of bool indicating if expert is enabled, and list of missing expert info
+    """
     missing = []
     provider = args.provider
     expert_provider = args.expert_provider
 
-    # Main provider keys
+    # Check API keys based on provider
     if provider == "anthropic":
         if not os.environ.get('ANTHROPIC_API_KEY'):
             missing.append('ANTHROPIC_API_KEY environment variable is not set')
@@ -278,63 +295,74 @@ def validate_environment(args):
         if not os.environ.get('OPENAI_API_BASE'):
             missing.append('OPENAI_API_BASE environment variable is not set')
 
-    # Expert keys
     expert_missing = []
+    # Check expert provider keys with fallback
     if expert_provider == "anthropic":
         expert_key_missing = not os.environ.get('EXPERT_ANTHROPIC_API_KEY')
         fallback_available = expert_provider == provider and os.environ.get('ANTHROPIC_API_KEY')
-        if expert_key_missing and not fallback_available:
+        if expert_key_missing and fallback_available:
+            os.environ['EXPERT_ANTHROPIC_API_KEY'] = os.environ['ANTHROPIC_API_KEY']
+            expert_key_missing = False
+        if expert_key_missing:
             expert_missing.append('EXPERT_ANTHROPIC_API_KEY environment variable is not set')
+
     elif expert_provider == "openai":
         expert_key_missing = not os.environ.get('EXPERT_OPENAI_API_KEY')
         fallback_available = expert_provider == provider and os.environ.get('OPENAI_API_KEY')
-        if expert_key_missing and not fallback_available:
+        if expert_key_missing and fallback_available:
+            os.environ['EXPERT_OPENAI_API_KEY'] = os.environ['OPENAI_API_KEY']
+            expert_key_missing = False
+        if expert_key_missing:
             expert_missing.append('EXPERT_OPENAI_API_KEY environment variable is not set')
+
     elif expert_provider == "openrouter":
         expert_key_missing = not os.environ.get('EXPERT_OPENROUTER_API_KEY')
         fallback_available = expert_provider == provider and os.environ.get('OPENROUTER_API_KEY')
-        if expert_key_missing and not fallback_available:
+        if expert_key_missing and fallback_available:
+            os.environ['EXPERT_OPENROUTER_API_KEY'] = os.environ['OPENROUTER_API_KEY']
+            expert_key_missing = False
+        if expert_key_missing:
             expert_missing.append('EXPERT_OPENROUTER_API_KEY environment variable is not set')
+
     elif expert_provider == "openai-compatible":
         expert_key_missing = not os.environ.get('EXPERT_OPENAI_API_KEY')
         fallback_available = expert_provider == provider and os.environ.get('OPENAI_API_KEY')
-        if expert_key_missing and not fallback_available:
+        if expert_key_missing and fallback_available:
+            os.environ['EXPERT_OPENAI_API_KEY'] = os.environ['OPENAI_API_KEY']
+            expert_key_missing = False
+        if expert_key_missing:
             expert_missing.append('EXPERT_OPENAI_API_KEY environment variable is not set')
+            
         expert_base_missing = not os.environ.get('EXPERT_OPENAI_API_BASE')
         base_fallback_available = expert_provider == provider and os.environ.get('OPENAI_API_BASE')
-        if expert_base_missing and not base_fallback_available:
+        if expert_base_missing and base_fallback_available:
+            os.environ['EXPERT_OPENAI_API_BASE'] = os.environ['OPENAI_API_BASE']
+            expert_base_missing = False
+        if expert_base_missing:
             expert_missing.append('EXPERT_OPENAI_API_BASE environment variable is not set')
 
-    # If main keys missing, just disable LLM entirely
-    llm_enabled = True
-    if missing:
-        llm_enabled = False
 
-    # If expert keys missing, disable expert
+    # If main keys missing, we must exit immediately
+    if missing:
+        print_error("Missing required dependencies:")
+        for item in missing:
+            print_error(f"- {item}")
+        sys.exit(1)
+
+    # If expert keys missing, we disable expert tools instead of exiting
     expert_enabled = True
     if expert_missing:
         expert_enabled = False
 
-    return llm_enabled, expert_enabled, missing, expert_missing
+    return expert_enabled, expert_missing
 
 def main():
+    """Main entry point for the ra-aid command line tool."""
     try:
         try:
             args = parse_arguments()
-            llm_enabled, expert_enabled, missing, expert_missing = validate_environment(args)
+            expert_enabled, expert_missing = validate_environment(args)  # Will exit if main env vars missing
             
-            # If main LLM keys missing
-            if missing:
-                # Disable LLM completely
-                console.print(Panel(
-                    f"[yellow]LLM disabled due to missing main provider configuration:[/yellow]\n" + 
-                    "\n".join(f"- {m}" for m in missing) +
-                    "\nSet the required environment variables to enable LLM features.",
-                    title="LLM Disabled",
-                    style="yellow"
-                ))
-            
-            # If expert keys missing
             if expert_missing:
                 console.print(Panel(
                     f"[yellow]Expert tools disabled due to missing configuration:[/yellow]\n" + 
@@ -344,7 +372,10 @@ def main():
                     style="yellow"
                 ))
             
-            # If no message, exit
+            # Create the base model after validation
+            model = initialize_llm(args.provider, args.model)
+
+            # Validate message is provided
             if not args.message:
                 print_error("--message is required")
                 sys.exit(1)
@@ -359,28 +390,20 @@ def main():
                 "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
             
-            # Only initialize the model if LLM is enabled
+            # Run research stage
-            model = None
-            if llm_enabled:
-                model = initialize_llm(args.provider, args.model)
-            
             print_stage_header("Research Stage")
             
-            research_tools = get_research_tools(
+            # Create research agent
-                research_only=_global_memory.get('config', {}).get('research_only', False),
-                expert_enabled=expert_enabled,
-                llm_enabled=llm_enabled
-            )
-            
-            # If no LLM, we can't run the agent, just print a message
-            if llm_enabled:
             research_agent = create_react_agent(
                 model,
-                    research_tools,
+                get_research_tools(research_only=_global_memory.get('config', {}).get('research_only', False), expert_enabled=expert_enabled),
                 checkpointer=research_memory
             )
             
@@ -389,22 +412,23 @@ def main():
 {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.'}"""
+
             try:
                 run_agent_with_retry(research_agent, research_prompt, config)
             except TaskCompletedException as e:
                 print_stage_header("Task Completed")
-                    raise
+                raise  # Re-raise to be caught by outer handler
-            else:
-                console.print(Panel("[yellow]LLM is disabled, cannot perform LLM-based research[/yellow]", title="No LLM Available"))
 
-            run_research_subtasks(base_task, config, model, expert_enabled=expert_enabled, llm_enabled=llm_enabled)
+            # 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")
-                planning_tools = get_planning_tools(expert_enabled=expert_enabled, llm_enabled=llm_enabled)
                 
-                if llm_enabled:
+                # Create planning agent
-                    planning_agent = create_react_agent(model, planning_tools, checkpointer=planning_memory)
+                planning_agent = create_react_agent(model, get_planning_tools(expert_enabled=expert_enabled), checkpointer=planning_memory)
+                
                 planning_prompt = PLANNING_PROMPT.format(
                     research_notes=get_memory_value('research_notes'),
                     key_facts=get_memory_value('key_facts'),
@@ -412,18 +436,18 @@ Be very thorough in your research and emit lots of snippets, key facts. If you t
                     base_task=base_task,
                     related_files="\n".join(get_related_files())
                 )
-                    run_agent_with_retry(planning_agent, planning_prompt, config)
-                else:
-                    console.print(Panel("[yellow]LLM is disabled, cannot perform planning[/yellow]", title="No LLM Available"))
 
+                # Run planning agent
+                run_agent_with_retry(planning_agent, planning_prompt, config)
+
+                # Run implementation stage with task-specific agents
                 run_implementation_stage(
                     base_task,
                     get_memory_value('tasks'),
                     get_memory_value('plan'),
                     get_related_files(),
                     model,
-                    expert_enabled=expert_enabled,
+                    expert_enabled=expert_enabled
-                    llm_enabled=llm_enabled
                 )
         except TaskCompletedException:
             sys.exit(0)