Add tests

Author: stellasia

src/neo4j_graphrag/llm/vertexai_llm.py

@@ -194,12 +194,12 @@ def _to_vertexai_tool(self, tool: Tool) -> VertexAITool:
                 FunctionDeclaration(
                     name=tool.get_name(),
                     description=tool.get_description(),
-                    parameters=tool.get_parameters(),
+                    parameters=tool.get_parameters(exclude=["additional_properties"]),
                 )
             ]
         )
 
-    def get_tools(
+    def _get_llm_tools(
         self, tools: Optional[Sequence[Tool]]
     ) -> Optional[list[VertexAITool]]:
         if not tools:
@@ -212,7 +212,7 @@ def _get_model(
         tools: Optional[Sequence[Tool]] = None,
     ) -> GenerativeModel:
         system_message = [system_instruction] if system_instruction is not None else []
-        vertex_ai_tools = self.get_tools(tools)
+        vertex_ai_tools = self._get_llm_tools(tools)
         model = GenerativeModel(
             model_name=self.model_name,
             system_instruction=system_message,
@@ -228,7 +228,7 @@ async def _acall_llm(
         system_instruction: Optional[str] = None,
         tools: Optional[Sequence[Tool]] = None,
     ) -> GenerationResponse:
-        model = self._get_model(system_instruction, tools)
+        model = self._get_model(system_instruction=system_instruction, tools=tools)
         messages = self.get_messages(input, message_history)
         response = await model.generate_content_async(messages, **self.model_params)
         return response
@@ -240,7 +240,7 @@ def _call_llm(
         system_instruction: Optional[str] = None,
         tools: Optional[Sequence[Tool]] = None,
     ) -> GenerationResponse:
-        model = self._get_model(system_instruction, tools)
+        model = self._get_model(system_instruction=system_instruction, tools=tools)
         messages = self.get_messages(input, message_history)
         response = model.generate_content(messages, **self.model_params)
         return response

src/neo4j_graphrag/tool.py

@@ -169,18 +169,21 @@ def _preprocess_properties(cls, values: dict[str, Any]) -> dict[str, Any]:
             values["properties"] = new_props
         return values
 
-    def model_dump_tool(self) -> Dict[str, Any]:
+    def model_dump_tool(self, exclude: Optional[list[str]] = None) -> Dict[str, Any]:
+        exclude = exclude or []
         properties_dict: Dict[str, Any] = {}
         for name, param in self.properties.items():
+            if name in exclude:
+                continue
             properties_dict[name] = param.model_dump_tool()
 
         result = super().model_dump_tool()
         result["properties"] = properties_dict
 
-        if self.required_properties:
+        if self.required_properties and "required" not in exclude:
             result["required"] = self.required_properties
 
-        if not self.additional_properties:
+        if not self.additional_properties and "additional_properties" not in exclude:
             result["additionalProperties"] = False
 
         return result
@@ -242,13 +245,13 @@ def get_description(self) -> str:
         """
         return self._description
 
-    def get_parameters(self) -> Dict[str, Any]:
+    def get_parameters(self, exclude: Optional[list[str]] = None) -> Dict[str, Any]:
         """Get the parameters the tool accepts in a dictionary format suitable for LLM providers.
 
         Returns:
             Dict[str, Any]: Dictionary containing parameter schema information.
         """
-        return self._parameters.model_dump_tool()
+        return self._parameters.model_dump_tool(exclude)
 
     def execute(self, query: str, **kwargs: Any) -> Any:
         """Execute the tool with the given query and additional parameters.

tests/unit/llm/conftest.py

@@ -0,0 +1,27 @@
+import pytest
+
+from neo4j_graphrag.tool import Tool, ObjectParameter, StringParameter
+
+
+class TestTool(Tool):
+    """Test tool for unit tests."""
+
+    def __init__(self, name: str = "test_tool", description: str = "A test tool"):
+        parameters = ObjectParameter(
+            description="Test parameters",
+            properties={"param1": StringParameter(description="Test parameter")},
+            required_properties=["param1"],
+            additional_properties=False,
+        )
+
+        super().__init__(
+            name=name,
+            description=description,
+            parameters=parameters,
+            execute_func=lambda **kwargs: kwargs,
+        )
+
+
+@pytest.fixture
+def test_tool() -> Tool:
+    return TestTool()

tests/unit/llm/test_openai_llm.py

@@ -20,7 +20,7 @@
 from neo4j_graphrag.llm import LLMResponse
 from neo4j_graphrag.llm.openai_llm import AzureOpenAILLM, OpenAILLM
 from neo4j_graphrag.llm.types import ToolCallResponse
-from neo4j_graphrag.tool import Tool, ObjectParameter, StringParameter
+from neo4j_graphrag.tool import Tool
 
 
 def get_mock_openai() -> MagicMock:
@@ -29,25 +29,6 @@ def get_mock_openai() -> MagicMock:
     return mock
 
 
-class TestTool(Tool):
-    """Test tool for unit tests."""
-
-    def __init__(self, name: str = "test_tool", description: str = "A test tool"):
-        parameters = ObjectParameter(
-            description="Test parameters",
-            properties={"param1": StringParameter(description="Test parameter")},
-            required_properties=["param1"],
-            additional_properties=False,
-        )
-
-        super().__init__(
-            name=name,
-            description=description,
-            parameters=parameters,
-            execute_func=lambda **kwargs: kwargs,
-        )
-
-
 @patch("builtins.__import__", side_effect=ImportError)
 def test_openai_llm_missing_dependency(mock_import: Mock) -> None:
     with pytest.raises(ImportError):
@@ -156,7 +137,9 @@ def test_openai_llm_with_message_history_validation_error(mock_import: Mock) ->
 @patch("builtins.__import__")
 @patch("json.loads")
 def test_openai_llm_invoke_with_tools_happy_path(
-    mock_json_loads: Mock, mock_import: Mock
+    mock_json_loads: Mock,
+    mock_import: Mock,
+    test_tool: Tool,
 ) -> None:
     # Set up json.loads to return a dictionary
     mock_json_loads.return_value = {"param1": "value1"}
@@ -183,7 +166,7 @@ def test_openai_llm_invoke_with_tools_happy_path(
     )
 
     llm = OpenAILLM(api_key="my key", model_name="gpt")
-    tools = [TestTool()]
+    tools = [test_tool]
 
     res = llm.invoke_with_tools("my text", tools)
     assert isinstance(res, ToolCallResponse)
@@ -196,7 +179,9 @@ def test_openai_llm_invoke_with_tools_happy_path(
 @patch("builtins.__import__")
 @patch("json.loads")
 def test_openai_llm_invoke_with_tools_with_message_history(
-    mock_json_loads: Mock, mock_import: Mock
+    mock_json_loads: Mock,
+    mock_import: Mock,
+    test_tool: Tool,
 ) -> None:
     # Set up json.loads to return a dictionary
     mock_json_loads.return_value = {"param1": "value1"}
@@ -223,7 +208,7 @@ def test_openai_llm_invoke_with_tools_with_message_history(
     )
 
     llm = OpenAILLM(api_key="my key", model_name="gpt")
-    tools = [TestTool()]
+    tools = [test_tool]
 
     message_history = [
         {"role": "user", "content": "When does the sun come up in the summer?"},
@@ -259,7 +244,9 @@ def test_openai_llm_invoke_with_tools_with_message_history(
 @patch("builtins.__import__")
 @patch("json.loads")
 def test_openai_llm_invoke_with_tools_with_system_instruction(
-    mock_json_loads: Mock, mock_import: Mock
+    mock_json_loads: Mock,
+    mock_import: Mock,
+    test_tool: Mock,
 ) -> None:
     # Set up json.loads to return a dictionary
     mock_json_loads.return_value = {"param1": "value1"}
@@ -286,7 +273,7 @@ def test_openai_llm_invoke_with_tools_with_system_instruction(
     )
 
     llm = OpenAILLM(api_key="my key", model_name="gpt")
-    tools = [TestTool()]
+    tools = [test_tool]
 
     system_instruction = "You are a helpful assistant."
 
@@ -314,7 +301,7 @@ def test_openai_llm_invoke_with_tools_with_system_instruction(
 
 
 @patch("builtins.__import__")
-def test_openai_llm_invoke_with_tools_error(mock_import: Mock) -> None:
+def test_openai_llm_invoke_with_tools_error(mock_import: Mock, test_tool: Tool) -> None:
     mock_openai = get_mock_openai()
     mock_import.return_value = mock_openai
 
@@ -324,7 +311,7 @@ def test_openai_llm_invoke_with_tools_error(mock_import: Mock) -> None:
     )
 
     llm = OpenAILLM(api_key="my key", model_name="gpt")
-    tools = [TestTool()]
+    tools = [test_tool]
 
     with pytest.raises(LLMGenerationError):
         llm.invoke_with_tools("my text", tools)

tests/unit/llm/test_vertexai_llm.py

@@ -19,9 +19,15 @@
 
 import pytest
 from neo4j_graphrag.exceptions import LLMGenerationError
+from neo4j_graphrag.llm.types import ToolCallResponse
 from neo4j_graphrag.llm.vertexai_llm import VertexAILLM
+from neo4j_graphrag.tool import Tool
 from neo4j_graphrag.types import LLMMessage
-from vertexai.generative_models import Content, Part
+from vertexai.generative_models import (
+    Content,
+    GenerationResponse,
+    Part,
+)
 
 
 @patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel", None)
@@ -171,4 +177,117 @@ async def test_vertexai_ainvoke_happy_path(GenerativeModelMock: MagicMock) -> No
     input_text = "may thy knife chip and shatter"
     response = await llm.ainvoke(input_text)
     assert response.content == "Return text"
-    llm.model.generate_content_async.assert_called_once_with([mock.ANY], **model_params)
+    llm.model.generate_content_async.assert_awaited_once_with(
+        [mock.ANY], **model_params
+    )
+
+
+def test_vertexai_get_llm_tools(test_tool: Tool) -> None:
+    llm = VertexAILLM(model_name="gemini")
+    tools = llm._get_llm_tools(tools=[test_tool])
+    assert tools is not None
+    assert len(tools) == 1
+    tool = tools[0]
+    tool_dict = tool.to_dict()
+    assert len(tool_dict["function_declarations"]) == 1
+    assert tool_dict["function_declarations"][0]["name"] == "test_tool"
+    assert tool_dict["function_declarations"][0]["description"] == "A test tool"
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._parse_tool_response")
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._call_llm")
+def test_vertexai_invoke_with_tools(
+    mock_call_llm: Mock, mock_parse_tool: Mock, test_tool: Tool,
+) -> None:
+    # Mock the model call response
+    tool_call_mock = MagicMock()
+    tool_call_mock.name = "function"
+    tool_call_mock.args = {}
+    mock_call_llm.return_value = MagicMock(
+        candidates=[MagicMock(function_calls=[tool_call_mock])]
+    )
+    mock_parse_tool.return_value = ToolCallResponse(tool_calls=[])
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    res = llm.invoke_with_tools("my text", tools)
+    mock_call_llm.assert_called_once_with(
+        "my text",
+        message_history=None,
+        system_instruction=None,
+        tools=tools,
+    )
+    mock_parse_tool.assert_called_once()
+    assert isinstance(res, ToolCallResponse)
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._get_model")
+def test_vertexai_call_llm_with_tools(mock_model: Mock, test_tool: Tool) -> None:
+    # Mock the generation response
+    mock_generate_content = mock_model.return_value.generate_content
+    mock_generate_content.return_value = MagicMock(
+        spec=GenerationResponse,
+    )
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    res = llm._call_llm("my text", tools=tools)
+    assert isinstance(res, GenerationResponse)
+
+    mock_model.assert_called_once_with(
+        system_instruction=None,
+        tools=tools,
+    )
+
+
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._parse_tool_response")
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._call_llm")
+def test_vertexai_ainvoke_with_tools(
+    mock_call_llm: Mock, mock_parse_tool: Mock, test_tool: Tool,
+) -> None:
+    # Mock the model call response
+    tool_call_mock = MagicMock()
+    tool_call_mock.name = "function"
+    tool_call_mock.args = {}
+    mock_call_llm.return_value = AsyncMock(
+        return_value=MagicMock(candidates=[MagicMock(function_calls=[tool_call_mock])])
+    )
+    mock_parse_tool.return_value = ToolCallResponse(tool_calls=[])
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    res = llm.invoke_with_tools("my text", tools)
+    mock_call_llm.assert_called_once_with(
+        "my text",
+        message_history=None,
+        system_instruction=None,
+        tools=tools,
+    )
+    mock_parse_tool.assert_called_once()
+    assert isinstance(res, ToolCallResponse)
+
+
+@pytest.mark.asyncio
+@patch("neo4j_graphrag.llm.vertexai_llm.VertexAILLM._get_model")
+async def test_vertexai_acall_llm_with_tools(mock_model, test_tool: Tool) -> None:
+    # Mock the generation response
+    mock_model.return_value = AsyncMock(
+        generate_content_async=AsyncMock(
+            return_value=MagicMock(
+                spec=GenerationResponse,
+            )
+        )
+    )
+
+    llm = VertexAILLM(model_name="gemini")
+    tools = [test_tool]
+
+    res = await llm._acall_llm("my text", tools=tools)
+    mock_model.assert_called_once_with(
+        system_instruction=None,
+        tools=tools,
+    )
+    assert isinstance(res, GenerationResponse)