Add Tool class

To not rely on json schema from openai

Author: oskarhane

examples/tool_calls/openai_tool_calls.py

@@ -17,32 +17,32 @@
 
 from neo4j_graphrag.llm import OpenAILLM
 from neo4j_graphrag.llm.types import ToolCallResponse
+from neo4j_graphrag.tool import Tool, ObjectParameter, StringParameter, IntegerParameter
 
 # Load environment variables from .env file
 load_dotenv()
 
-# Define a tool for extracting information from text
-TOOLS = [
-    {
-        "type": "function",
-        "function": {
-            "name": "extract_person_info",
-            "description": "Extract information about a person from text",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "name": {"type": "string", "description": "The person's full name"},
-                    "age": {"type": "integer", "description": "The person's age"},
-                    "occupation": {
-                        "type": "string",
-                        "description": "The person's occupation",
-                    },
-                },
-                "required": ["name"],
-            },
-        },
-    }
-]
+
+# Create a custom Tool implementation for person info extraction
+parameters = ObjectParameter(
+    description="Parameters for extracting person information",
+    properties={
+        "name": StringParameter(description="The person's full name"),
+        "age": IntegerParameter(description="The person's age"),
+        "occupation": StringParameter(description="The person's occupation"),
+    },
+    required_properties=["name"],
+    additional_properties=False,
+)
+person_info_tool = Tool(
+    name="extract_person_info",
+    description="Extract information about a person from text",
+    parameters=parameters,
+    execute_func=lambda **kwargs: kwargs,
+)
+
+# Create the tool instance
+TOOLS = [person_info_tool]
 
 
 def process_tool_call(response: ToolCallResponse) -> Dict[str, Any]:

src/neo4j_graphrag/tool.py

@@ -204,6 +204,37 @@ def validate_properties(self) -> "ObjectParameter":
         self.properties = validated_properties
         return self
 
+class ArrayParameter(ToolParameter):
+    """Array parameter for tools."""
+
+    def __init__(
+        self,
+        description: str,
+        items: ToolParameter,
+        required: bool = False,
+        min_items: Optional[int] = None,
+        max_items: Optional[int] = None,
+    ):
+        super().__init__(description, required)
+        self.items = items
+        self.min_items = min_items
+        self.max_items = max_items
+
+    def to_dict(self) -> Dict[str, Any]:
+        result: Dict[str, Any] = {
+            "type": ParameterType.ARRAY,
+            "description": self.description,
+            "items": self.items.to_dict(),
+        }
+
+        if self.min_items is not None:
+            result["minItems"] = self.min_items
+
+        if self.max_items is not None:
+            result["maxItems"] = self.max_items
+
+        return result
+
 
 class Tool(ABC):
     """Abstract base class defining the interface for all tools in the neo4j-graphrag library."""