Coverage for agentos/core/handoff.py: 53%

1""" 

2Handoff protocol for NexusAgent. 

3 

4Provides Swarm-style task transfer between agents. 

5When an agent cannot handle a request, it can transfer 

6to another agent that is better suited. 

7""" 

8 

9from __future__ import annotations 

10 

11from dataclasses import dataclass, field 

12from typing import Any, TypeVar 

13 

14from agentos.core.di import Agent, RunContext 

15 

16# Type variable for agent 

17T = TypeVar("T") 

18 

19 

20@dataclass 

21class Handoff: 

22 """ 

23 Represents a handoff request to another agent. 

24 

25 Usage: 

26 class SupportAgent(Agent[str, str]): 

27 async def run(self, ctx: RunContext[str]) -> str | Handoff: 

28 if "billing" in ctx.deps.lower(): 

29 return transfer_to(BillingAgent(), ctx.deps) 

30 return "General support" 

31 """ 

32 

33 target_agent: Agent[Any, Any] 

34 input_data: Any 

35 metadata: dict[str, Any] = field(default_factory=dict) 

36 reason: str = "" 

37 

38 def __post_init__(self): 

39 """Validate handoff.""" 

40 if self.target_agent is None: 

41 raise ValueError("target_agent cannot be None") 

42 

43 

44@dataclass 

45class HandoffResult: 

46 """ 

47 Result of a handoff operation. 

48 

49 Contains: 

50 - output: The final output from the target agent 

51 - source_agent: Name of the original agent 

52 - target_agent: Name of the agent that handled it 

53 - handoff_chain: List of agents involved 

54 """ 

55 

56 output: Any 

57 source_agent: str 

58 target_agent: str 

59 handoff_chain: list[str] = field(default_factory=list) 

60 metadata: dict[str, Any] = field(default_factory=dict) 

61 

62 

63def transfer_to(agent: Agent[Any, Any], input_data: Any, reason: str = "", **metadata) -> Handoff: 

64 """ 

65 Create a handoff to another agent. 

66 

67 Args: 

68 agent: Target agent to transfer to 

69 input_data: Data to pass to the target agent 

70 reason: Reason for the handoff 

71 **metadata: Additional metadata 

72 

73 Returns: 

74 Handoff object 

75 

76 Usage: 

77 return transfer_to(BillingAgent(), ctx.deps, reason="Billing question") 

78 """ 

79 return Handoff( 

80 target_agent=agent, 

81 input_data=input_data, 

82 metadata=metadata, 

83 reason=reason, 

84 ) 

85 

86 

87def can_handle(agent: Agent[Any, Any], input_data: Any) -> bool: 

88 """ 

89 Check if an agent can handle the input. 

90 

91 This is a helper function that calls the agent's 

92 can_handle() method if it exists, otherwise returns True. 

93 

94 Args: 

95 agent: Agent to check 

96 input_data: Input data 

97 

98 Returns: 

99 True if agent can handle, False otherwise 

100 """ 

101 if hasattr(agent, "can_handle"): 

102 return agent.can_handle(input_data) 

103 return True 

104 

105 

106async def execute_with_handoff( 

107 agent: Agent[Any, Any], input_data: Any, max_hops: int = 10, **metadata 

108) -> HandoffResult | Any: 

109 """ 

110 Execute an agent with automatic handoff handling. 

111 

112 If the agent returns a Handoff, automatically execute 

113 the target agent and return the result. 

114 

115 Args: 

116 agent: Starting agent 

117 input_data: Input data 

118 max_hops: Maximum number of handoffs 

119 **metadata: Additional metadata 

120 

121 Returns: 

122 HandoffResult if handoffs occurred, otherwise raw output 

123 

124 Raises: 

125 RuntimeError: If max_hops exceeded 

126 """ 

127 current_agent = agent 

128 current_input = input_data 

129 handoff_chain = [current_agent.name] 

130 

131 for hop in range(max_hops): 

132 # Execute current agent 

133 result = await current_agent.invoke(current_input, **metadata) 

134 

135 # Check if result is a handoff 

136 if isinstance(result, Handoff): 

137 # Move to next agent 

138 current_agent = result.target_agent 

139 current_input = result.input_data 

140 handoff_chain.append(current_agent.name) 

141 

142 # Merge metadata 

143 metadata.update(result.metadata) 

144 else: 

145 # No handoff, we're done 

146 if len(handoff_chain) > 1: 

147 # Return HandoffResult if we had handoffs 

148 return HandoffResult( 

149 output=result, 

150 source_agent=handoff_chain[0], 

151 target_agent=handoff_chain[-1], 

152 handoff_chain=handoff_chain, 

153 metadata=metadata, 

154 ) 

155 else: 

156 # No handoffs, return raw output 

157 return result 

158 

159 raise RuntimeError(f"Max handoff hops ({max_hops}) exceeded") 

160 

161 

162class HandoffAwareAgent(Agent[Any, Any]): 

163 """ 

164 Base class for agents that support handoffs. 

165 

166 Provides can_handle() method for checking if agent 

167 can handle input, and run() can return Handoff. 

168 """ 

169 

170 def can_handle(self, input_data: Any) -> bool: 

171 """ 

172 Check if this agent can handle the input. 

173 

174 Override in subclass to add custom logic. 

175 

176 Args: 

177 input_data: Input data 

178 

179 Returns: 

180 True if can handle, False otherwise 

181 """ 

182 return True 

183 

184 async def run(self, ctx: RunContext[Any]) -> Any: 

185 """ 

186 Main agent logic. Can return Handoff to transfer. 

187 

188 Override in subclass. 

189 """ 

190 raise NotImplementedError("Subclass must implement run()") 

191 

192 

193# ── Auto-generated compat stubs ──