Adding new description generator

This commit is contained in:
João Moura
2024-02-29 03:09:48 -03:00
parent 640b5a9461
commit ec97e15a3a
19 changed files with 53 additions and 10 deletions

View File

@@ -13,10 +13,12 @@ class BaseTool(BaseModel, ABC):
"""Used to tell the model how/when/why to use the tool."""
args_schema: Optional[Type[V1BaseModel]] = None
"""The schema for the arguments that the tool accepts."""
description_updated: bool = False
@model_validator(mode="after")
def _check_args_schema(self):
self._set_args_schema()
self._generate_description()
return self
def run(
@@ -56,6 +58,13 @@ class BaseTool(BaseModel, ABC):
},
},
)
def _generate_description(self):
args = []
for arg, attribute in self.args_schema.schema()['properties'].items():
args.append(f"{arg}: '{attribute['type']}'")
description = self.description.replace('\n', ' ')
self.description = f"{self.name}({', '.join(args)}) - {description}"
class Tool(BaseTool):