added custom tools

Author: dillonalaird

tests/tools/test_tools.py

@@ -2,8 +2,10 @@
 import tempfile
 
 import numpy as np
+import pytest
 from PIL import Image
 
+from vision_agent.tools import TOOLS, Tool, register_tool
 from vision_agent.tools.tools import BboxIoU, BoxDistance, SegArea, SegIoU
 
 
@@ -65,3 +67,71 @@ def test_box_distance():
     box1 = [0, 0, 2, 2]
     box2 = [1, 1, 3, 3]
     assert box_dist(box1, box2) == 0.0
+
+
+def test_register_tool():
+    assert TOOLS[len(TOOLS) - 1]["name"] != "test_tool_"
+
+    @register_tool
+    class TestTool(Tool):
+        name = "test_tool_"
+        description = "Test Tool"
+        usage = {
+            "required_parameters": [{"name": "prompt", "type": "str"}],
+            "examples": [
+                {
+                    "scenario": "Test",
+                    "parameters": {"prompt": "Test Prompt"},
+                }
+            ],
+        }
+
+        def __call__(self, prompt: str) -> str:
+            return prompt
+
+    assert TOOLS[len(TOOLS) - 1]["name"] == "test_tool_"
+
+
+def test_register_tool_incorrect():
+    with pytest.raises(ValueError):
+
+        @register_tool
+        class NoAttributes(Tool):
+            pass
+
+    with pytest.raises(ValueError):
+
+        @register_tool
+        class NoName(Tool):
+            description = "Test Tool"
+            usage = {
+                "required_parameters": [{"name": "prompt", "type": "str"}],
+                "examples": [
+                    {
+                        "scenario": "Test",
+                        "parameters": {"prompt": "Test Prompt"},
+                    }
+                ],
+            }
+
+    with pytest.raises(ValueError):
+
+        @register_tool
+        class NoDescription(Tool):
+            name = "test_tool_"
+            usage = {
+                "required_parameters": [{"name": "prompt", "type": "str"}],
+                "examples": [
+                    {
+                        "scenario": "Test",
+                        "parameters": {"prompt": "Test Prompt"},
+                    }
+                ],
+            }
+
+    with pytest.raises(ValueError):
+
+        @register_tool
+        class NoUsage(Tool):
+            name = "test_tool_"
+            description = "Test Tool"

vision_agent/tools/__init__.py

@@ -14,4 +14,5 @@
     SegArea,
     SegIoU,
     Tool,
+    register_tool,
 )

vision_agent/tools/tools.py

@@ -2,7 +2,7 @@
 import tempfile
 from abc import ABC
 from pathlib import Path
-from typing import Any, Dict, List, Tuple, Union, cast
+from typing import Any, Dict, List, Tuple, Type, Union, cast
 
 import numpy as np
 import requests
@@ -765,6 +765,30 @@ def __call__(self, equation: str) -> float:
 }
 
 
+def register_tool(tool: Type[Tool]) -> None:
+    r"""Add a tool to the list of available tools.
+
+    Parameters:
+        tool: The tool to add.
+    """
+
+    if (
+        not hasattr(tool, "name")
+        or not hasattr(tool, "description")
+        or not hasattr(tool, "usage")
+    ):
+        raise ValueError(
+            "The tool must have 'name', 'description' and 'usage' attributes."
+        )
+
+    TOOLS[len(TOOLS)] = {
+        "name": tool.name,
+        "description": tool.description,
+        "usage": tool.usage,
+        "class": tool,
+    }
+
+
 def _send_inference_request(
     payload: Dict[str, Any], endpoint_name: str
 ) -> Dict[str, Any]: