|
|
|
|
|
import inspect |
|
|
from typing import Dict, List, Optional, Any |
|
|
|
|
|
from ..core.module import BaseModule |
|
|
|
|
|
ALLOWED_TYPES = ["string", "number", "integer", "boolean", "object", "array"] |
|
|
|
|
|
|
|
|
class Tool(BaseModule): |
|
|
name: str |
|
|
description: str |
|
|
inputs: Dict[str, Dict[str, Any]] |
|
|
required: Optional[List[str]] = None |
|
|
|
|
|
""" |
|
|
inputs: {"input_name": {"type": "string", "description": "input description"}, ...} |
|
|
""" |
|
|
|
|
|
def __init_subclass__(cls): |
|
|
super().__init_subclass__() |
|
|
cls.validate_attributes() |
|
|
|
|
|
def get_tool_schema(self) -> Dict: |
|
|
return { |
|
|
"type": "function", |
|
|
"function": { |
|
|
"name": self.name, |
|
|
"description": self.description, |
|
|
"parameters": { |
|
|
"type": "object", |
|
|
"properties": self.inputs, |
|
|
"required": self.required |
|
|
} |
|
|
} |
|
|
} |
|
|
|
|
|
@classmethod |
|
|
def validate_attributes(cls): |
|
|
required_attributes = { |
|
|
"name": str, |
|
|
"description": str, |
|
|
"inputs": dict |
|
|
} |
|
|
|
|
|
json_to_python = { |
|
|
"string": str, |
|
|
"integer": int, |
|
|
"number": float, |
|
|
"boolean": bool, |
|
|
"object": dict, |
|
|
"array": list, |
|
|
} |
|
|
|
|
|
for attr, attr_type in required_attributes.items(): |
|
|
if not hasattr(cls, attr): |
|
|
raise ValueError(f"Attribute {attr} is required") |
|
|
if not isinstance(getattr(cls, attr), attr_type): |
|
|
raise ValueError(f"Attribute {attr} must be of type {attr_type}") |
|
|
|
|
|
for input_name, input_content in cls.inputs.items(): |
|
|
if not isinstance(input_content, dict): |
|
|
raise ValueError(f"Input '{input_name}' must be a dictionary") |
|
|
if "type" not in input_content or "description" not in input_content: |
|
|
raise ValueError(f"Input '{input_name}' must have 'type' and 'description'") |
|
|
if input_content["type"] not in ALLOWED_TYPES: |
|
|
raise ValueError(f"Input '{input_name}' must have a valid type, should be one of {ALLOWED_TYPES}") |
|
|
|
|
|
call_signature = inspect.signature(cls.__call__) |
|
|
if input_name not in call_signature.parameters: |
|
|
raise ValueError(f"Input '{input_name}' is not found in __call__") |
|
|
if call_signature.parameters[input_name].annotation != json_to_python[input_content["type"]]: |
|
|
raise ValueError(f"Input '{input_name}' has a type mismatch in __call__") |
|
|
|
|
|
if cls.required: |
|
|
for required_input in cls.required: |
|
|
if required_input not in cls.inputs: |
|
|
raise ValueError(f"Required input '{required_input}' is not found in inputs") |
|
|
|
|
|
def __call__(self, **kwargs): |
|
|
raise NotImplementedError("All tools must implement __call__") |
|
|
|
|
|
class Toolkit(BaseModule): |
|
|
name: str |
|
|
tools: List[Tool] |
|
|
|
|
|
def get_tool_names(self) -> List[str]: |
|
|
return [tool.name for tool in self.tools] |
|
|
|
|
|
def get_tool_descriptions(self) -> List[str]: |
|
|
return [tool.description for tool in self.tools] |
|
|
|
|
|
def get_tool_inputs(self) -> List[Dict]: |
|
|
return [tool.inputs for tool in self.tools] |
|
|
|
|
|
def add_tool(self, tool: Tool): |
|
|
self.tools.append(tool) |
|
|
|
|
|
def remove_tool(self, tool_name: str): |
|
|
self.tools = [tool for tool in self.tools if tool.name != tool_name] |
|
|
|
|
|
def get_tool(self, tool_name: str) -> Tool: |
|
|
for tool in self.tools: |
|
|
if tool.name == tool_name: |
|
|
return tool |
|
|
raise ValueError(f"Tool '{tool_name}' not found") |
|
|
|
|
|
def get_tools(self) -> List[Tool]: |
|
|
return self.tools |
|
|
|
|
|
def get_tool_schemas(self) -> List[Dict]: |
|
|
return [tool.get_tool_schema() for tool in self.tools] |
|
|
|
