Coverage for src/command_registry.py: 82%

50 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 23:16 -0700

1""" 

2Unified command registry for both user commands and agent tools. 

3 

4This module provides a central registry for all commands that can be executed 

5by both the user interface and LLM agent tools, reducing code duplication. 

6""" 

7 

8from typing import Dict, Any, Callable, List, Optional 

9from dataclasses import dataclass, field 

10 

11 

12@dataclass 

13class CommandDefinition: 

14 """ 

15 Definition for a command that can be used by both TUI and agent. 

16 

17 Attributes: 

18 name: Command name (without leading slash) 

19 description: Human-readable description of the command 

20 handler: Function that implements the command logic 

21 parameters: JSON Schema compatible parameter definitions 

22 required_params: List of required parameter names 

23 needs_api_client: Whether the command needs an API client to execute 

24 visible_to_user: Whether the command should be visible in user interface 

25 visible_to_agent: Whether the command should be available to the agent 

26 tui_aliases: Alternative command names in the TUI (with leading slash) 

27 output_formatter: Optional custom output formatter function 

28 usage_hint: Optional[str] = None (for better TUI error messages) 

29 supports_interactive_input: bool = False (to make interactive mode explicit) 

30 agent_display: str = "condensed" (how to display when called by agent: "condensed" or "full") 

31 condensed_action: Optional[str] = None (friendly action name for condensed display, e.g. "Setting catalog") 

32 """ 

33 

34 name: str 

35 description: str 

36 handler: Callable 

37 parameters: Dict[str, Dict[str, Any]] = field(default_factory=dict) 

38 required_params: List[str] = field(default_factory=list) 

39 needs_api_client: bool = True 

40 visible_to_user: bool = True # Some commands may be agent-only 

41 visible_to_agent: bool = True # Some commands may be user-only 

42 tui_aliases: List[str] = field( 

43 default_factory=list 

44 ) # TUI alternatives like "/models" for "list_models" 

45 output_formatter: Optional[Callable] = None # Custom output formatter if needed 

46 usage_hint: Optional[str] = None # For better TUI error messages 

47 supports_interactive_input: bool = False # To make interactive mode explicit 

48 agent_display: str = ( 

49 "condensed" # How to display when called by agent: "condensed" or "full" 

50 ) 

51 condensed_action: Optional[str] = None # Friendly action name for condensed display 

52 

53 

54# Command registry - populated with all available commands 

55COMMAND_REGISTRY: Dict[str, CommandDefinition] = {} 

56TUI_COMMAND_MAP: Dict[str, str] = ( 

57 {} 

58) # Maps TUI command names (with slash) to registry names 

59 

60 

61def register_command(command_def: CommandDefinition) -> None: 

62 """ 

63 Register a command in the unified registry. 

64 

65 Args: 

66 command_def: Command definition to register 

67 """ 

68 COMMAND_REGISTRY[command_def.name] = command_def 

69 

70 # Also register TUI command aliases if any 

71 for alias in command_def.tui_aliases: 

72 if alias.startswith("/"): 

73 TUI_COMMAND_MAP[alias] = command_def.name 

74 else: 

75 TUI_COMMAND_MAP[f"/{alias}"] = command_def.name 

76 

77 

78def get_command(name: str) -> Optional[CommandDefinition]: 

79 """ 

80 Get a command by name. 

81 

82 Args: 

83 name: Command name to look up (without leading slash) 

84 

85 Returns: 

86 CommandDefinition if found, None otherwise 

87 """ 

88 # Direct lookup in registry 

89 if name in COMMAND_REGISTRY: 

90 return COMMAND_REGISTRY[name] 

91 

92 # Handle TUI command names (with leading slash) 

93 if name.startswith("/"): 

94 tui_name = name 

95 else: 

96 tui_name = f"/{name}" 

97 

98 # Check TUI command map 

99 if tui_name in TUI_COMMAND_MAP: 

100 registry_name = TUI_COMMAND_MAP[tui_name] 

101 return COMMAND_REGISTRY.get(registry_name) 

102 

103 return None 

104 

105 

106def get_user_commands() -> Dict[str, CommandDefinition]: 

107 """ 

108 Get all commands available to the user. 

109 

110 Returns: 

111 Dict mapping command names to definitions for user-visible commands 

112 """ 

113 return {name: cmd for name, cmd in COMMAND_REGISTRY.items() if cmd.visible_to_user} 

114 

115 

116def get_agent_commands() -> Dict[str, CommandDefinition]: 

117 """ 

118 Get all commands available to the agent. 

119 

120 Returns: 

121 Dict mapping command names to definitions for agent-visible commands 

122 """ 

123 return {name: cmd for name, cmd in COMMAND_REGISTRY.items() if cmd.visible_to_agent} 

124 

125 

126def get_agent_tool_schemas() -> List[Dict[str, Any]]: 

127 """ 

128 Get all command schemas in agent tool format. 

129 

130 Returns: 

131 List of tool schema definitions compatible with the LLM agent API 

132 """ 

133 agent_tools = [] 

134 for name, cmd in COMMAND_REGISTRY.items(): 

135 if not cmd.visible_to_agent: 

136 continue 

137 

138 tool = { 

139 "type": "function", 

140 "function": { 

141 "name": name, 

142 "description": cmd.description, 

143 "parameters": { 

144 "type": "object", 

145 "properties": cmd.parameters or {}, 

146 "required": cmd.required_params or [], 

147 }, 

148 }, 

149 } 

150 agent_tools.append(tool) 

151 return agent_tools 

152 

153 

154def resolve_tui_command(command: str) -> Optional[str]: 

155 """ 

156 Resolve a TUI command (with slash) to its registry name. 

157 

158 Args: 

159 command: TUI command (with leading slash) 

160 

161 Returns: 

162 Registry command name if found, None otherwise 

163 """ 

164 return TUI_COMMAND_MAP.get(command)