Coverage for src / osiris_cli / tools_registry.py: 0%

1#!/usr/bin/env python3 

2""" 

3Tool Registry - Centralized tool management 

4Handles tool definitions, execution, and permissions. 

5""" 

6 

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

8from dataclasses import dataclass 

9import inspect 

10 

11from .session import session 

12 

13 

14@dataclass 

15class ToolDefinition: 

16 """Represents a tool that can be called by the agent""" 

17 name: str 

18 description: str 

19 parameters: Dict[str, Any] 

20 function: Callable 

21 requires_approval: bool = False 

22 category: str = "general" 

23 

24 

25class ToolsRegistry: 

26 """ 

27 Central registry for all agent tools. 

28 Manages tool definitions, execution, and permissions. 

29 """ 

30 

31 def __init__(self): 

32 self._tools: Dict[str, ToolDefinition] = {} 

33 self._auto_approved_tools: set = set() 

34 

35 def register( 

36 self, 

37 name: str, 

38 description: str, 

39 parameters: Dict[str, Any], 

40 function: Callable, 

41 requires_approval: bool = False, 

42 category: str = "general" 

43 ): 

44 """Register a new tool""" 

45 self._tools[name] = ToolDefinition( 

46 name=name, 

47 description=description, 

48 parameters=parameters, 

49 function=function, 

50 requires_approval=requires_approval, 

51 category=category 

52 ) 

53 

54 def register_decorator( 

55 self, 

56 description: str, 

57 parameters: Dict[str, Any], 

58 requires_approval: bool = False, 

59 category: str = "general" 

60 ): 

61 """Decorator for registering tools""" 

62 def decorator(func: Callable): 

63 self.register( 

64 name=func.__name__, 

65 description=description, 

66 parameters=parameters, 

67 function=func, 

68 requires_approval=requires_approval, 

69 category=category 

70 ) 

71 return func 

72 return decorator 

73 

74 async def execute(self, name: str, arguments: Dict[str, Any]) -> Any: 

75 """Execute a tool by name (Asynchronous & Non-blocking)""" 

76 import asyncio 

77 if name not in self._tools: 

78 raise ValueError(f"Tool '{name}' not found") 

79 

80 tool = self._tools[name] 

81 

82 session.enforce_tool_policy(name) 

83 

84 # Check if tool requires approval and isn't auto-approved 

85 if tool.requires_approval and name not in self._auto_approved_tools: 

86 raise PermissionError(f"Tool '{name}' requires approval") 

87 

88 # Execute the tool safely 

89 if asyncio.iscoroutinefunction(tool.function): 

90 return await tool.function(**arguments) 

91 

92 # 2025 Best Practice: Run blocking sync tools in a separate thread 

93 # This ensures the TUI animation remains fluid. 

94 return await asyncio.to_thread(tool.function, **arguments) 

95 

96 def get_definitions(self) -> Dict: 

97 """Return tools as dict for display""" 

98 return { 

99 name: { 

100 "description": tool.description, 

101 "parameters": tool.parameters, 

102 } 

103 for name, tool in self._tools.items() 

104 } 

105 

106 def get_openai_format(self) -> List[Dict[str, Any]]: 

107 """Get tools in OpenAI function calling format""" 

108 return self.get_openai_format_filtered() 

109 

110 def get_openai_format_filtered(self, categories: Optional[List[str]] = None) -> List[Dict[str, Any]]: 

111 """Get filtered tools in OpenAI format. Always includes 'general'.""" 

112 allowed = set(categories) if categories else None 

113 

114 filtered = [] 

115 for tool in self._tools.values(): 

116 if allowed is None or tool.category == "general" or tool.category in allowed: 

117 filtered.append({ 

118 "type": "function", 

119 "function": { 

120 "name": tool.name, 

121 "description": tool.description, 

122 "parameters": tool.parameters 

123 } 

124 }) 

125 return filtered 

126 

127 def get_tool(self, name: str) -> Optional[ToolDefinition]: 

128 """Get a tool definition by name""" 

129 return self._tools.get(name) 

130 

131 def list_tools(self, category: Optional[str] = None) -> List[ToolDefinition]: 

132 """List all tools, optionally filtered by category""" 

133 if category: 

134 return [t for t in self._tools.values() if t.category == category] 

135 return list(self._tools.values()) 

136 

137 def auto_approve(self, tool_name: str): 

138 """Mark a tool as auto-approved""" 

139 self._auto_approved_tools.add(tool_name) 

140 

141 def load_skills_from_dir(self, skills_dir: str): 

142 """Dynamically load and register tools from a directory of Python files.""" 

143 import importlib.util 

144 import sys 

145 from pathlib import Path 

146 

147 path = Path(skills_dir) 

148 if not path.exists(): 

149 return 

150 

151 for py_file in path.glob("*.py"): 

152 if py_file.name.startswith("_"): continue 

153 

154 try: 

155 module_name = f"osiris_skill_{py_file.stem}" 

156 spec = importlib.util.spec_from_file_location(module_name, py_file) 

157 if spec and spec.loader: 

158 module = importlib.util.module_from_spec(spec) 

159 sys.modules[module_name] = module 

160 spec.loader.exec_module(module) 

161 

162 # Look for functions decorated with @tools_registry.register_decorator 

163 # or a specific 'register_skill' call in the module 

164 if hasattr(module, "register_skill"): 

165 module.register_skill(self) 

166 except Exception as e: 

167 print(f"Error loading skill {py_file.name}: {e}") 

168 

169 def remove_auto_approval(self, tool_name: str): 

170 """Remove auto-approval for a tool""" 

171 self._auto_approved_tools.discard(tool_name) 

172 

173 def is_auto_approved(self, tool_name: str) -> bool: 

174 """Check if a tool is auto-approved""" 

175 return tool_name in self._auto_approved_tools 

176 

177 

178# Global registry instance 

179tools_registry = ToolsRegistry()