Update dependencies (#384)

* Update deps

* Update lock file after merge

* mypy part 1

* Ruff

* mypy part 2

* Fix tests and update example

Author: stellasia

examples/customize/llms/vertexai_llm.py

@@ -1,12 +1,15 @@
 from neo4j_graphrag.llm import LLMResponse, VertexAILLM
 from vertexai.generative_models import GenerationConfig
 
-generation_config = GenerationConfig(temperature=0.0)
+generation_config = GenerationConfig(temperature=1.0)
 llm = VertexAILLM(
-    model_name="gemini-1.5-flash-001",
+    model_name="gemini-2.0-flash-001",
     generation_config=generation_config,
     # add here any argument that will be passed to the
     # vertexai.generative_models.GenerativeModel client
 )
-res: LLMResponse = llm.invoke("say something")
+res: LLMResponse = llm.invoke(
+    "say something",
+    system_instruction="You are living in 3000 where AI rules the world",
+)
 print(res.content)

poetry.lock

@@ -33,7 +33,7 @@ neo4j = "^5.17.0"
 pydantic = "^2.6.3"
 fsspec = "^2024.9.0"
 pypdf = "^5.1.0"
-json-repair = "^0.39.1"
+json-repair = "^0.44.1"
 pyyaml = "^6.0.2"
 types-pyyaml = "^6.0.12.20240917"
 # optional deps

pyproject.toml

@@ -21,7 +21,7 @@
 try:
     import vertexai
 except ImportError:
-    vertexai = None
+    vertexai = None  # type: ignore[assignment]
 
 
 class VertexAIEmbeddings(Embedder):

src/neo4j_graphrag/embeddings/vertexai.py

@@ -18,7 +18,7 @@
 import enum
 import json
 import logging
-from typing import Any, List, Optional, Union, cast
+from typing import Any, List, Optional, Union
 
 import json_repair
 from pydantic import ValidationError, validate_call
@@ -102,7 +102,7 @@ def balance_curly_braces(json_string: str) -> str:
 
 def fix_invalid_json(raw_json: str) -> str:
     repaired_json = json_repair.repair_json(raw_json)
-    repaired_json = cast(str, repaired_json).strip()
+    repaired_json = repaired_json.strip()
 
     if repaired_json == '""':
         raise InvalidJSONError("JSON repair resulted in an empty or invalid JSON.")

src/neo4j_graphrag/experimental/components/entity_relation_extractor.py

@@ -48,8 +48,8 @@
         ToolConfig,
     )
 except ImportError:
-    GenerativeModel = None
-    ResponseValidationError = None
+    GenerativeModel = None  # type: ignore[misc, assignment]
+    ResponseValidationError = None  # type: ignore[misc, assignment]
 
 
 class VertexAILLM(LLMInterface):
@@ -114,13 +114,15 @@ def get_messages(
                 if message.get("role") == "user":
                     messages.append(
                         Content(
-                            role="user", parts=[Part.from_text(message.get("content"))]
+                            role="user",
+                            parts=[Part.from_text(message.get("content", ""))],
                         )
                     )
                 elif message.get("role") == "assistant":
                     messages.append(
                         Content(
-                            role="model", parts=[Part.from_text(message.get("content"))]
+                            role="model",
+                            parts=[Part.from_text(message.get("content", ""))],
                         )
                     )
 
@@ -211,10 +213,10 @@ def _get_model(
         self,
         system_instruction: Optional[str] = None,
     ) -> GenerativeModel:
-        system_message = [system_instruction] if system_instruction is not None else []
+        # system_message = [system_instruction] if system_instruction is not None else []
         model = GenerativeModel(
             model_name=self.model_name,
-            system_instruction=system_message,
+            system_instruction=system_instruction,
         )
         return model
 
@@ -253,7 +255,7 @@ async def _acall_llm(
         model = self._get_model(system_instruction=system_instruction)
         options = self._get_call_params(input, message_history, tools)
         response = await model.generate_content_async(**options)
-        return response
+        return response  # type: ignore[no-any-return]
 
     def _call_llm(
         self,
@@ -265,7 +267,7 @@ def _call_llm(
         model = self._get_model(system_instruction=system_instruction)
         options = self._get_call_params(input, message_history, tools)
         response = model.generate_content(**options)
-        return response
+        return response  # type: ignore[no-any-return]
 
     def _to_tool_call(self, function_call: FunctionCall) -> ToolCall:
         return ToolCall(

src/neo4j_graphrag/llm/vertexai_llm.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
+from typing import Any
 from unittest import mock
 from unittest.mock import MagicMock, Mock
 
@@ -235,7 +236,7 @@ async def test_run(_: Mock, driver: MagicMock) -> None:
         },
         database_=None,
     )
-    parameters_ = {
+    parameters_: dict[str, Any] = {
         "rows": [
             {
                 "type": "RELATIONSHIP",
@@ -293,7 +294,7 @@ async def test_run_is_version_below_5_23(_: Mock) -> None:
         },
         database_=None,
     )
-    parameters_ = {
+    parameters_: dict[str, Any] = {
         "rows": [
             {
                 "type": "RELATIONSHIP",
@@ -352,7 +353,7 @@ async def test_run_is_version_5_23_or_above(_: Mock) -> None:
         },
         database_=None,
     )
-    parameters_ = {
+    parameters_: dict[str, Any] = {
         "rows": [
             {
                 "type": "RELATIONSHIP",

tests/unit/experimental/components/test_kg_writer.py

@@ -50,7 +50,7 @@ def test_vertexai_invoke_happy_path(GenerativeModelMock: MagicMock) -> None:
     assert response.content == "Return text"
     GenerativeModelMock.assert_called_once_with(
         model_name=model_name,
-        system_instruction=[],
+        system_instruction=None,
     )
     last_call = mock_model.generate_content.call_args_list[0]
     content = last_call.kwargs["contents"]
@@ -82,7 +82,7 @@ def test_vertexai_invoke_with_system_instruction(
     assert response.content == "Return text"
     GenerativeModelMock.assert_called_once_with(
         model_name=model_name,
-        system_instruction=[system_instruction],
+        system_instruction=system_instruction,
     )
     mock_model.generate_content.assert_called_once_with(
         contents=[{"text": "some text"}]
@@ -116,7 +116,7 @@ def test_vertexai_invoke_with_message_history_and_system_instruction(
     assert response.content == "Return text"
     GenerativeModelMock.assert_called_once_with(
         model_name=model_name,
-        system_instruction=[system_instruction],
+        system_instruction=system_instruction,
     )
     last_call = mock_model.generate_content.call_args_list[0]
     content = last_call.kwargs["contents"]