allow select tool at role initialization & restructure writecodewithtools

This commit is contained in:
yzlin 2024-01-20 21:06:48 +08:00
parent 2ccfe31123
commit 540542834e
8 changed files with 127 additions and 90 deletions

View file

@ -11,6 +11,7 @@ import re
from collections import defaultdict
import yaml
from pydantic import BaseModel
from metagpt.const import TOOL_SCHEMA_PATH
from metagpt.logs import logger
@ -18,11 +19,10 @@ from metagpt.tools.tool_convert import convert_code_to_tool_schema
from metagpt.tools.tool_data_type import Tool, ToolSchema, ToolType
class ToolRegistry:
def __init__(self):
self.tools = {}
self.tool_types = {}
self.tools_by_types = defaultdict(dict) # two-layer k-v, {tool_type: {tool_name: {...}, ...}, ...}
class ToolRegistry(BaseModel):
tools: dict = {}
tool_types: dict = {}
tools_by_types: dict = defaultdict(dict) # two-layer k-v, {tool_type: {tool_name: {...}, ...}, ...}
def register_tool_type(self, tool_type: ToolType):
self.tool_types[tool_type.name] = tool_type
@ -70,22 +70,22 @@ class ToolRegistry:
self.tools_by_types[tool_type][tool_name] = tool
logger.info(f"{tool_name} registered")
def has_tool(self, key):
def has_tool(self, key: str) -> Tool:
return key in self.tools
def get_tool(self, key):
def get_tool(self, key) -> Tool:
return self.tools.get(key)
def get_tools_by_type(self, key):
return self.tools_by_types.get(key)
def get_tools_by_type(self, key) -> dict[str, Tool]:
return self.tools_by_types.get(key, {})
def has_tool_type(self, key):
def has_tool_type(self, key) -> bool:
return key in self.tool_types
def get_tool_type(self, key):
def get_tool_type(self, key) -> ToolType:
return self.tool_types.get(key)
def get_tool_types(self):
def get_tool_types(self) -> dict[str, ToolType]:
return self.tool_types
@ -141,3 +141,16 @@ def make_schema(tool_source_object, include, path):
print(e)
return schema
def validate_tool_names(tools: list[str], return_tool_object=False) -> list[str]:
valid_tools = []
for tool_name in tools:
if not TOOL_REGISTRY.has_tool(tool_name):
logger.warning(
f"Specified tool {tool_name} not found and was skipped. Check if you have registered it properly"
)
else:
valid_tool = TOOL_REGISTRY.get_tool(tool_name) if return_tool_object else tool_name
valid_tools.append(valid_tool)
return valid_tools