fix tool useage

FyZyX · Jun 18, 2023 · 0c624b6 · 0c624b6
1 parent 7c5b616
commit 0c624b6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="llmfuncs",
-    version="0.2.1",
+    version="0.2.2",
     url="https://github.com/FyZyX/llmfuncs",
     author="Lucas Lofaro",
     author_email="[email protected]",

diff --git a/src/llmfuncs/tool.py b/src/llmfuncs/tool.py
@@ -64,7 +64,7 @@ def schema(self):
 
 class ToolCollection:
     def __init__(self, tools: typing.List[Tool] = None):
-        self._tools = {}
+        self._tools: typing.Dict[str, Tool] = {}
         for tool in tools or []:
             self.add_tool(tool)
 
@@ -128,12 +128,11 @@ def use_tool(self, tool_name: str, json_args: str | typing.Mapping) -> typing.An
         if not tool:
             raise ValueError(f"No tool found with name: {tool_name}")
 
-        func = tool['function']
-        params_schema = tool['schema']['parameters']
+        params_schema = tool.schema()['parameters']
         is_string = isinstance(json_args, str)
         args = validator.parse_json(json_args) if is_string else json_args
         validator.validate_args_with_schema(args, params_schema)
-        return func(**args)
+        return tool(**args)
 
     def schema(self) -> typing.List[schema.JsonSchema]:
         return [tool.schema() for tool in self._tools.values()]