Tribh/devops-copilot

1from typing import Any, Callable, Dict, List, Optional, Type2from pydantic import BaseModel, ConfigDict, Field, create_model, validate_call3import inspect4from devops_copilot.utils.logger import logger5 6class Tool(BaseModel):7    """Represents a tool available to agents."""8    name: str9    description: str10    parameters_schema: Type[BaseModel]11    func: Callable12    13    model_config = ConfigDict(arbitrary_types_allowed=True)14 15    def execute(self, **kwargs) -> Any:16        """17        Executes the tool with validation.18        IMPORTANT: All arguments are validated against parameters_schema before execution.19        """20        logger.info(f"Executing tool: {self.name} with params: {kwargs}")21        validated_params = self.parameters_schema(**kwargs)22        return self.func(**validated_params.model_dump())23 24class ToolRegistry:25    """Registry for managing and validating tools."""26    27    def __init__(self):28        self._tools: Dict[str, Tool] = {}29 30    def register(self, name: str, description: str):31        """Decorator to register a function as a tool."""32        def decorator(func: Callable):33            # Create a Pydantic model from the function signature34            sig = inspect.signature(func)35            fields = {}36            for param_name, param in sig.parameters.items():37                if param.annotation == inspect.Parameter.empty:38                    fields[param_name] = (Any, ...)39                else:40                    fields[param_name] = (param.annotation, ... if param.default == inspect.Parameter.empty else param.default)41            42            pydantic_model = create_model(f"{func.__name__}_Schema", **fields)43            44            tool = Tool(45                name=name,46                description=description,47                parameters_schema=pydantic_model,48                func=func49            )50            self._tools[name] = tool51            logger.info(f"Registered tool: {name}")52            return func53        return decorator54 55    def get_tool(self, name: str) -> Optional[Tool]:56        return self._tools.get(name)57 58    def list_tools(self) -> List[Dict[str, Any]]:59        return [60            {61                "name": tool.name,62                "description": tool.description,63                "parameters": tool.parameters_schema.model_json_schema()64            }65            for tool in self._tools.values()66        ]67 68# Global registry instance69registry = ToolRegistry()70