diff --git a/src/cai/cli.py b/src/cai/cli.py index 5182b494..6c2c3646 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -202,6 +202,57 @@ agent = Agent( ) ) +def update_agent_models_recursively(agent, new_model, visited=None): + """ + Recursively update the model for an agent and all agents in its handoffs. + + Args: + agent: The agent to update + new_model: The new model string to set + visited: Set of agent names already visited to prevent infinite loops + """ + if visited is None: + visited = set() + + # Avoid infinite loops by tracking visited agents + if agent.name in visited: + return + visited.add(agent.name) + + # Update the main agent's model + if hasattr(agent, 'model') and hasattr(agent.model, 'model'): + agent.model.model = new_model + + # Update models for all handoff agents + if hasattr(agent, 'handoffs'): + for handoff_item in agent.handoffs: + # Handle both direct Agent references and Handoff objects + if hasattr(handoff_item, 'on_invoke_handoff'): + # This is a Handoff object + # For handoffs created with the handoff() function, the agent is stored + # in the closure of the on_invoke_handoff function + # We can try to extract it from the function's closure + try: + # Get the closure variables of the handoff function + if hasattr(handoff_item.on_invoke_handoff, '__closure__') and handoff_item.on_invoke_handoff.__closure__: + for cell in handoff_item.on_invoke_handoff.__closure__: + if hasattr(cell.cell_contents, 'model') and hasattr(cell.cell_contents, 'name'): + # This looks like an agent + handoff_agent = cell.cell_contents + update_agent_models_recursively( + handoff_agent, new_model, visited + ) + break + except Exception: + # If we can't extract the agent from closure, skip it + pass + elif hasattr(handoff_item, 'model'): + # This is a direct Agent reference + update_agent_models_recursively( + handoff_item, new_model, visited + ) + + def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), force_until_flag=False): """ Run a simple interactive CLI loop for CAI. @@ -302,10 +353,9 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), # Check if model has changed and update if needed current_model = os.getenv('CAI_MODEL', 'qwen2.5:14b') if current_model != last_model and hasattr(agent, 'model'): - # Update the model in the agent - if hasattr(agent.model, 'model'): - agent.model.model = current_model - last_model = current_model + # Update the model recursively for the agent and all handoff agents + update_agent_models_recursively(agent, current_model) + last_model = current_model # Check if agent type has changed and recreate agent if needed current_agent_type = os.getenv('CAI_AGENT_TYPE', 'one_tool_agent') @@ -322,9 +372,8 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), if hasattr(agent.model, 'suppress_final_output'): agent.model.suppress_final_output = False # Changed to False to show all agent messages - # Apply current model to the new agent - if hasattr(agent.model, 'model'): - agent.model.model = current_model + # Apply current model to the new agent and all its handoff agents + update_agent_models_recursively(agent, current_model) # Set agent name in the model for streaming display if hasattr(agent.model, 'set_agent_name'): @@ -540,10 +589,9 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), # Create a fresh agent instance instance_agent = get_agent_by_name(config.agent_name) - # Override model if specified + # Override model if specified, updating recursively if config.model and hasattr(instance_agent, 'model'): - if hasattr(instance_agent.model, 'model'): - instance_agent.model.model = config.model + update_agent_models_recursively(instance_agent, config.model) # Override prompt if specified instance_input = config.prompt or input_text @@ -665,7 +713,7 @@ def run_cai_cli(starting_agent, context_variables=None, max_turns=float('inf'), # Configure agent instance to match main agent settings if hasattr(instance_agent, 'model') and hasattr(agent, 'model'): if hasattr(instance_agent.model, 'model') and hasattr(agent.model, 'model'): - instance_agent.model.model = agent.model.model + update_agent_models_recursively(instance_agent, agent.model.model) # Create a fresh input for this instance - use just the user input directly # This ensures each instance has its own completely independent context @@ -877,6 +925,10 @@ def main(): if hasattr(agent.model, 'suppress_final_output'): agent.model.suppress_final_output = False # Changed to False to show all agent messages + # Ensure the agent and all its handoff agents use the current model + current_model = os.getenv('CAI_MODEL', 'qwen2.5:14b') + update_agent_models_recursively(agent, current_model) + # Run the CLI with the selected agent run_cai_cli(agent)