Coverage for core / src / sensorkit / cli / agent.py: 0%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncclick as click 

3 

4from sensorkit.auto.agent import AgentControllerInfo, AgentState 

5from sensorkit.cli.utils import entity_option, with_kit 

6 

7 

8@click.group("agent") 

9async def agent_group(): 

10 """Agent management commands.""" 

11 pass 

12 

13 

14@agent_group.command("include") 

15@click.argument("program_name") 

16@entity_option(default="agent", help="Agent entity name") 

17@with_kit 

18async def include(kit, program_name: str, entity: str): 

19 """Include a program in the agent's management.""" 

20 from sensorkit.auto.agent import AgentConfigureRequest, agent_configure_request 

21 

22 await kit.entity(entity).call( 

23 agent_configure_request, 

24 AgentConfigureRequest(remove_program_exclusions={program_name}), 

25 ) 

26 click.echo(f"Removed scheduling exclusion for '{program_name}'") 

27 

28 

29@agent_group.command("exclude") 

30@click.argument("program_name") 

31@entity_option(default="agent", help="Agent entity name") 

32@with_kit 

33async def exclude(kit, program_name: str, entity: str): 

34 """Exclude a program from the agent's management.""" 

35 from sensorkit.auto.agent import AgentConfigureRequest, agent_configure_request 

36 

37 await kit.entity(entity).call( 

38 agent_configure_request, 

39 AgentConfigureRequest(add_program_exclusions={program_name}), 

40 ) 

41 click.echo(f"Added scheduling exclusion for '{program_name}'") 

42 

43 

44@agent_group.command("global-control") 

45@click.argument("state", type=click.Choice(["on", "off"], case_sensitive=False)) 

46@entity_option(default="agent", help="Agent entity name") 

47@with_kit 

48async def global_control(kit, state: str, entity: str): 

49 """Enable or disable global control.""" 

50 from sensorkit.auto.agent import AgentConfigureRequest, agent_configure_request 

51 

52 enabled = state.lower() == "on" 

53 await kit.entity(entity).call( 

54 agent_configure_request, 

55 AgentConfigureRequest(global_control_enabled=enabled), 

56 ) 

57 click.echo(f"Global control was set to {state}") 

58 

59 

60@agent_group.command("control") 

61@click.argument("controller") 

62@click.argument("state", type=click.Choice(["on", "off"], case_sensitive=False)) 

63@entity_option(default="agent", help="Agent entity name") 

64@with_kit 

65async def controller_control(kit, controller: str, state: str, entity: str): 

66 """Enable or disable control for a specific controller.""" 

67 from sensorkit.auto.agent import AgentConfigureRequest, agent_configure_request 

68 

69 enabled = state.lower() == "on" 

70 await kit.entity(entity).call( 

71 agent_configure_request, 

72 AgentConfigureRequest(controller_control_enabled={controller: enabled}), 

73 ) 

74 click.echo(f"Control for '{controller}' was set to {state}") 

75 

76 

77@agent_group.command("override") 

78@click.argument("controller") 

79@click.argument("state", type=click.Choice(["up", "down", "none"], case_sensitive=False)) 

80@entity_option(default="agent", help="Agent entity name") 

81@with_kit 

82async def controller_override(kit, controller: str, state: str, entity: str): 

83 """Set or clear demand override for a specific controller.""" 

84 from sensorkit.auto.agent import AgentConfigureRequest, agent_configure_request 

85 

86 state = state.lower() 

87 override = True if state == "up" else False if state == "down" else None 

88 

89 await kit.entity(entity).call( 

90 agent_configure_request, 

91 AgentConfigureRequest(controller_demand_override={controller: override}), 

92 ) 

93 click.echo(f"Demand override for '{controller}' was set to {state.upper()}") 

94 

95 

96@agent_group.command("scheduling") 

97@click.argument("state", type=click.Choice(["on", "off"], case_sensitive=False)) 

98@entity_option(default="agent", help="Agent entity name") 

99@with_kit 

100async def scheduling(kit, state: str, entity: str): 

101 """Enable or disable scheduling.""" 

102 from sensorkit.auto.agent import AgentConfigureRequest, agent_configure_request 

103 

104 enabled = state.lower() == "on" 

105 await kit.entity(entity).call( 

106 agent_configure_request, AgentConfigureRequest(scheduling_enabled=enabled) 

107 ) 

108 click.echo(f"Scheduling was set to {state}") 

109 

110 

111@agent_group.command("status") 

112@entity_option(default="agent", help="Agent entity name") 

113@with_kit 

114async def status(kit, entity: str): 

115 """View the agent's current status and configuration.""" 

116 from sensorkit.auto.agent import AgentState 

117 

118 try: 

119 state = await kit.entity(entity).kv_get_model(AgentState) 

120 except Exception as e: 

121 click.echo(f"Error: Could not recover agent state: {e}", err=True) 

122 return 

123 

124 click.echo(f"{'Agent:':<18} {entity}") 

125 

126 # Global settings 

127 click.echo( 

128 f"{'Global Control:':<18} {'ON' if state.operating_state.global_control_enabled else 'OFF'}" 

129 ) 

130 click.echo( 

131 f"{'Scheduling:':<18} {'ON' if state.scheduler_state.scheduling_enabled else 'OFF'}\n" 

132 ) 

133 

134 if state.operating_state.controllers: 

135 for i, (name, info) in enumerate(state.operating_state.controllers.items(), 1): 

136 _print_controller_info(i, name, info, state) 

137 click.echo() 

138 else: 

139 click.echo("No controllers registered in agent.\n") 

140 

141 # Excluded programs 

142 if state.scheduler_state.excluded_programs: 

143 click.echo("Excluded Programs:") 

144 

145 for program in sorted(state.scheduler_state.excluded_programs): 

146 click.echo(f" - {program}") 

147 else: 

148 click.echo("No excluded programs.") 

149 

150 

151def _print_controller_info(i: int, name: str, info: AgentControllerInfo, state: AgentState): 

152 control = "ON" if info.control_enabled else "OFF" 

153 

154 if not state.operating_state.global_control_enabled: 

155 control += "*" 

156 

157 override = "NONE" 

158 

159 if info.demand_override is True: 

160 override = "UP" 

161 elif info.demand_override is False: 

162 override = "DOWN" 

163 

164 elected = "NONE" 

165 

166 if info.elected_state is True: 

167 elected = "UP" 

168 elif info.elected_state is False: 

169 elected = "DOWN" 

170 

171 click.echo(f"{f'Controller {i}:':<18} {name}\n") 

172 click.echo(f" Control: {control}") 

173 click.echo(f" State override: {override}") 

174 click.echo(f" Elected state: {elected}") 

175 click.echo() 

176 

177 if schedule := state.scheduler_state.schedule.get(name): 

178 from intervaltree import IntervalTree 

179 

180 from sensorkit.auto.scheduler import debug_print_schedule 

181 

182 debug_print_schedule( 

183 IntervalTree.from_tuples(schedule), 

184 lookahead_mins=5, 

185 print_func=lambda x: click.echo(f" {x}"), 

186 ) 

187 else: 

188 click.echo(" No schedule present.")