Use pydantic model for the tool parameters

Author: oskarhane

src/neo4j_graphrag/tool.py

@@ -1,6 +1,7 @@
 from abc import ABC
 from enum import Enum
-from typing import Any, Dict, List, Callable, Optional
+from typing import Any, Dict, List, Callable, Optional, Union, ClassVar, Type
+from pydantic import BaseModel, Field, model_validator
 
 
 class ParameterType(str, Enum):
@@ -14,56 +15,61 @@ class ParameterType(str, Enum):
     ARRAY = "array"
 
 
-class ToolParameter:
-    """Base class for all tool parameters."""
-
-    def __init__(self, description: str, required: bool = False):
-        self.description = description
-        self.required = required
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert the parameter to a dictionary format."""
-        raise NotImplementedError("Subclasses must implement to_dict")
+class ToolParameter(BaseModel):
+    """Base class for all tool parameters using Pydantic."""
+    description: str
+    required: bool = False
+    type: ClassVar[ParameterType]
+    
+    def model_dump_tool(self) -> Dict[str, Any]:
+        """Convert the parameter to a dictionary format for tool usage."""
+        result = {"type": self.type, "description": self.description}
+        return result
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ToolParameter":
+        """Create a parameter from a dictionary."""
+        param_type = data.get("type")
+        if not param_type:
+            raise ValueError("Parameter type is required")
+            
+        # Find the appropriate class based on the type
+        param_classes = {
+            ParameterType.STRING: StringParameter,
+            ParameterType.INTEGER: IntegerParameter,
+            ParameterType.NUMBER: NumberParameter,
+            ParameterType.BOOLEAN: BooleanParameter,
+            ParameterType.OBJECT: ObjectParameter,
+            ParameterType.ARRAY: ArrayParameter,
+        }
+        
+        param_class = param_classes.get(param_type)
+        if not param_class:
+            raise ValueError(f"Unknown parameter type: {param_type}")
+            
+        return param_class.model_validate(data)
 
 
 class StringParameter(ToolParameter):
     """String parameter for tools."""
-
-    def __init__(
-        self, description: str, required: bool = False, enum: Optional[List[str]] = None
-    ):
-        super().__init__(description, required)
-        self.enum = enum
-
-    def to_dict(self) -> Dict[str, Any]:
-        result: Dict[str, Any] = {
-            "type": ParameterType.STRING,
-            "description": self.description,
-        }
+    type: ClassVar[ParameterType] = ParameterType.STRING
+    enum: Optional[List[str]] = None
+    
+    def model_dump_tool(self) -> Dict[str, Any]:
+        result = super().model_dump_tool()
         if self.enum:
             result["enum"] = self.enum
         return result
 
 
 class IntegerParameter(ToolParameter):
     """Integer parameter for tools."""
-
-    def __init__(
-        self,
-        description: str,
-        required: bool = False,
-        minimum: Optional[int] = None,
-        maximum: Optional[int] = None,
-    ):
-        super().__init__(description, required)
-        self.minimum = minimum
-        self.maximum = maximum
-
-    def to_dict(self) -> Dict[str, Any]:
-        result: Dict[str, Any] = {
-            "type": ParameterType.INTEGER,
-            "description": self.description,
-        }
+    type: ClassVar[ParameterType] = ParameterType.INTEGER
+    minimum: Optional[int] = None
+    maximum: Optional[int] = None
+    
+    def model_dump_tool(self) -> Dict[str, Any]:
+        result = super().model_dump_tool()
         if self.minimum is not None:
             result["minimum"] = self.minimum
         if self.maximum is not None:
@@ -73,23 +79,12 @@ def to_dict(self) -> Dict[str, Any]:
 
 class NumberParameter(ToolParameter):
     """Number parameter for tools."""
-
-    def __init__(
-        self,
-        description: str,
-        required: bool = False,
-        minimum: Optional[float] = None,
-        maximum: Optional[float] = None,
-    ):
-        super().__init__(description, required)
-        self.minimum = minimum
-        self.maximum = maximum
-
-    def to_dict(self) -> Dict[str, Any]:
-        result: Dict[str, Any] = {
-            "type": ParameterType.NUMBER,
-            "description": self.description,
-        }
+    type: ClassVar[ParameterType] = ParameterType.NUMBER
+    minimum: Optional[float] = None
+    maximum: Optional[float] = None
+    
+    def model_dump_tool(self) -> Dict[str, Any]:
+        result = super().model_dump_tool()
         if self.minimum is not None:
             result["minimum"] = self.minimum
         if self.maximum is not None:
@@ -99,77 +94,71 @@ def to_dict(self) -> Dict[str, Any]:
 
 class BooleanParameter(ToolParameter):
     """Boolean parameter for tools."""
+    type: ClassVar[ParameterType] = ParameterType.BOOLEAN
+
 
-    def to_dict(self) -> Dict[str, Any]:
-        return {"type": ParameterType.BOOLEAN, "description": self.description}
+class ArrayParameter(ToolParameter):
+    """Array parameter for tools."""
+    type: ClassVar[ParameterType] = ParameterType.ARRAY
+    items: "ToolParameter"
+    min_items: Optional[int] = None
+    max_items: Optional[int] = None
+    
+    def model_dump_tool(self) -> Dict[str, Any]:
+        result = super().model_dump_tool()
+        result["items"] = self.items.model_dump_tool()
+        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
+    
+    @model_validator(mode="after")
+    def validate_items(self) -> "ArrayParameter":
+        if not isinstance(self.items, ToolParameter):
+            if isinstance(self.items, dict):
+                self.items = ToolParameter.from_dict(self.items)
+            else:
+                raise ValueError(f"Items must be a ToolParameter or dict, got {type(self.items)}")
+        return self
 
 
 class ObjectParameter(ToolParameter):
     """Object parameter for tools."""
-
-    def __init__(
-        self,
-        description: str,
-        properties: Dict[str, ToolParameter],
-        required: bool = False,
-        required_properties: Optional[List[str]] = None,
-        additional_properties: bool = True,
-    ):
-        super().__init__(description, required)
-        self.properties = properties
-        self.required_properties = required_properties or []
-        self.additional_properties = additional_properties
-
-    def to_dict(self) -> Dict[str, Any]:
+    type: ClassVar[ParameterType] = ParameterType.OBJECT
+    properties: Dict[str, ToolParameter]
+    required_properties: List[str] = Field(default_factory=list)
+    additional_properties: bool = True
+    
+    def model_dump_tool(self) -> Dict[str, Any]:
         properties_dict: Dict[str, Any] = {}
         for name, param in self.properties.items():
-            properties_dict[name] = param.to_dict()
-
-        result: Dict[str, Any] = {
-            "type": ParameterType.OBJECT,
-            "description": self.description,
-            "properties": properties_dict,
-        }
+            properties_dict[name] = param.model_dump_tool()
 
+        result = super().model_dump_tool()
+        result["properties"] = properties_dict
+        
         if self.required_properties:
             result["required"] = self.required_properties
-
+            
         if not self.additional_properties:
             result["additionalProperties"] = False
-
-        return result
-
-
-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
+    
+    @model_validator(mode="after")
+    def validate_properties(self) -> "ObjectParameter":
+        validated_properties = {}
+        for name, param in self.properties.items():
+            if not isinstance(param, ToolParameter):
+                if isinstance(param, dict):
+                    validated_properties[name] = ToolParameter.from_dict(param)
+                else:
+                    raise ValueError(f"Property {name} must be a ToolParameter or dict, got {type(param)}")
+            else:
+                validated_properties[name] = param
+        self.properties = validated_properties
+        return self
 
 
 class Tool(ABC):
@@ -179,12 +168,18 @@ def __init__(
         self,
         name: str,
         description: str,
-        parameters: ObjectParameter,
+        parameters: Union[ObjectParameter, Dict[str, Any]],
         execute_func: Callable[..., Any],
     ):
         self._name = name
         self._description = description
-        self._parameters = parameters
+        
+        # Allow parameters to be provided as a dictionary
+        if isinstance(parameters, dict):
+            self._parameters = ObjectParameter.model_validate(parameters)
+        else:
+            self._parameters = parameters
+            
         self._execute_func = execute_func
 
     def get_name(self) -> str:
@@ -209,7 +204,7 @@ def get_parameters(self) -> Dict[str, Any]:
         Returns:
             Dict[str, Any]: Dictionary containing parameter schema information.
         """
-        return self._parameters.to_dict()
+        return self._parameters.model_dump_tool()
 
     def execute(self, query: str, **kwargs: Any) -> Any:
         """Execute the tool with the given query and additional parameters.