Add "required" parameter to ToolParameter.model_dump_tool

Author: oskarhane

src/neo4j_graphrag/tool.py

@@ -24,7 +24,9 @@ class ToolParameter(BaseModel):
 
     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}
+        result: Dict[str, Any] = {"type": self.type, "description": self.description}
+        if self.required:
+            result["required"] = True
         return result
 
     @classmethod

tests/unit/tool/test_tool.py

@@ -21,6 +21,7 @@ def test_string_parameter() -> None:
     d = param.model_dump_tool()
     assert d["type"] == ParameterType.STRING
     assert d["enum"] == ["a", "b"]
+    assert d["required"] is True
 
 
 def test_integer_parameter() -> None:
@@ -139,6 +140,39 @@ def test_from_dict() -> None:
         ToolParameter.from_dict({"description": "no type"})
 
 
+def test_required_parameter() -> None:
+    # Test that required=True is included in model_dump_tool output for different parameter types
+    string_param = StringParameter(description="Required string", required=True)
+    assert string_param.model_dump_tool()["required"] is True
+
+    integer_param = IntegerParameter(description="Required integer", required=True)
+    assert integer_param.model_dump_tool()["required"] is True
+
+    number_param = NumberParameter(description="Required number", required=True)
+    assert number_param.model_dump_tool()["required"] is True
+
+    boolean_param = BooleanParameter(description="Required boolean", required=True)
+    assert boolean_param.model_dump_tool()["required"] is True
+
+    array_param = ArrayParameter(
+        description="Required array",
+        items=StringParameter(description="item"),
+        required=True,
+    )
+    assert array_param.model_dump_tool()["required"] is True
+
+    object_param = ObjectParameter(
+        description="Required object",
+        properties={"prop": StringParameter(description="property")},
+        required=True,
+    )
+    assert object_param.model_dump_tool()["required"] is True
+
+    # Test that required=False doesn't include the required field
+    optional_param = StringParameter(description="Optional string", required=False)
+    assert "required" not in optional_param.model_dump_tool()
+
+
 def test_tool_class() -> None:
     def dummy_func(query: str, **kwargs: Any) -> dict[str, Any]:
         return kwargs