Make Neo4j version check work with trailing "-XXX" (#162)

* Fix version verification

* CI

* Ruff

Author: stellasia

src/neo4j_graphrag/retrievers/base.py

@@ -87,6 +87,20 @@ def __init__(self, driver: neo4j.Driver, neo4j_database: Optional[str] = None):
         if self.VERIFY_NEO4J_VERSION:
             self._verify_version()
 
+    def _get_version(self) -> tuple[tuple[int, ...], bool]:
+        records, _, _ = self.driver.execute_query(
+            "CALL dbms.components()", database_=self.neo4j_database
+        )
+        version = records[0]["versions"][0]
+        # drop everything after the '-' first
+        version_main, *_ = version.split("-")
+        # convert each number between '.' into int
+        version_tuple = tuple(map(int, version_main.split(".")))
+        # if no patch version, consider it's 0
+        if len(version_tuple) < 3:
+            version_tuple = (*version_tuple, 0)
+        return version_tuple, "aura" in version
+
     def _verify_version(self) -> None:
         """
         Check if the connected Neo4j database version supports vector indexing.
@@ -96,19 +110,11 @@ def _verify_version(self) -> None:
         indexing. Raises a Neo4jMinVersionError if the connected Neo4j version is
         not supported.
         """
-        records, _, _ = self.driver.execute_query(
-            "CALL dbms.components()", database_=self.neo4j_database
-        )
-        version = records[0]["versions"][0]
+        version_tuple, is_aura = self._get_version()
 
-        if "aura" in version:
-            version_tuple = (
-                *tuple(map(int, version.split("-")[0].split("."))),
-                0,
-            )
+        if is_aura:
             target_version = (5, 18, 0)
         else:
-            version_tuple = tuple(map(int, version.split(".")))
             target_version = (5, 18, 1)
 
         if version_tuple < target_version:

tests/unit/retrievers/test_base.py

@@ -24,25 +24,54 @@
 from neo4j_graphrag.types import RawSearchResult, RetrieverResult
 
 
+@pytest.mark.parametrize(
+    "db_version,expected_version",
+    [
+        (["5.18-aura"], ((5, 18, 0), True)),
+        (["5.3-aura"], ((5, 3, 0), True)),
+        (["5.19.0"], ((5, 19, 0), False)),
+        (["4.3.5"], ((4, 3, 5), False)),
+        (["5.23.0-6698"], ((5, 23, 0), False)),
+    ],
+)
+def test_retriever_get_version(
+    driver: MagicMock,
+    db_version: list[str],
+    expected_version: tuple[tuple[int, ...], bool],
+) -> None:
+    class MockRetriever(Retriever):
+        VERIFY_NEO4J_VERSION = False
+
+        def get_search_results(self, *args: Any, **kwargs: Any) -> RawSearchResult:
+            return RawSearchResult(records=[])
+
+    driver.execute_query.return_value = [[{"versions": db_version}], None, None]
+    retriever = MockRetriever(driver)
+    assert retriever._get_version() == expected_version
+
+
 @pytest.mark.parametrize(
     "db_version,expected_exception",
     [
-        (["5.18-aura"], None),
-        (["5.3-aura"], Neo4jVersionError),
-        (["5.19.0"], None),
-        (["4.3.5"], Neo4jVersionError),
+        (((5, 18, 0), True), None),
+        (((5, 3, 0), True), Neo4jVersionError),
+        (((5, 19, 0), False), None),
+        (((4, 3, 5), False), Neo4jVersionError),
+        (((5, 23, 0), False), None),
     ],
 )
+@patch("neo4j_graphrag.retrievers.base.Retriever._get_version")
 def test_retriever_version_support(
+    mock_get_version: MagicMock,
     driver: MagicMock,
-    db_version: list[str],
+    db_version: tuple[tuple[int, ...], bool],
     expected_exception: Union[type[ValueError], None],
 ) -> None:
     class MockRetriever(Retriever):
         def get_search_results(self, *args: Any, **kwargs: Any) -> RawSearchResult:
             return RawSearchResult(records=[])
 
-    driver.execute_query.return_value = [[{"versions": db_version}], None, None]
+    mock_get_version.return_value = db_version
     if expected_exception:
         with pytest.raises(expected_exception):
             MockRetriever(driver=driver)