Rewrite model._model_info to use call_endpoint

Author: FlorentinD

graphdatascience/model/link_prediction_model.py

@@ -41,5 +41,5 @@ def link_features(self) -> List[LinkFeature]:
             A list of LinkFeatures of the pipeline.
 
         """
-        steps: List[Dict[str, Any]] = self._list_info()["modelInfo"][0]["pipeline"]["featureSteps"]
+        steps: List[Dict[str, Any]] = self._list_info()["modelInfo"]["pipeline"]["featureSteps"]
         return [LinkFeature(s["name"], s["config"]) for s in steps]

graphdatascience/model/model.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 from pandas import DataFrame, Series
 
@@ -23,33 +23,30 @@ def __init__(self, name: str, query_runner: QueryRunner, server_version: ServerV
     def _endpoint_prefix(self) -> str:
         pass
 
-    def _list_info(self) -> DataFrame:
+    def _list_info(self) -> Series[Any]:
+        params = CallParameters(name=self.name())
+
+        result: Series[Any]
         if self._server_version < ServerVersion(2, 5, 0):
-            query = "CALL gds.beta.model.list($name)"
+            result = self._query_runner.call_procedure(
+                "gds.beta.model.list", params=params, custom_error=False
+            ).squeeze()
         else:
-            query = """
-                    CALL gds.model.list($name)
-                    YIELD
-                      modelName, modelType, modelInfo,
-                      creationTime, trainConfig, graphSchema,
-                      loaded, stored, published
-                    RETURN
-                      modelName, modelType,
-                      modelInfo {.*, modelName: modelName, modelType: modelType} AS modelInfo,
-                      creationTime, trainConfig, graphSchema,
-                      loaded, stored, published, published AS shared
-                    """
-
-        params = {"name": self.name()}
-
-        # FIXME use call procedure + do post processing on the client side
-        # TODO really fixme
-        info = self._query_runner.run_cypher(query, params, custom_error=False)
-
-        if len(info) == 0:
+            result = self._query_runner.call_procedure("gds.model.list", params=params, custom_error=False).squeeze()
+
+            if not result.empty:
+                #  2.5 compat format
+                result["modelInfo"] = {
+                    **result["modelInfo"],
+                    "modelName": result["modelName"],
+                    "modelType": result["modelType"],
+                }
+                result["shared"] = result["published"]
+
+        if result.empty:
             raise ValueError(f"There is no '{self.name()}' in the model catalog")
 
-        return info
+        return result
 
     def _estimate_predict(self, predict_mode: str, graph_name: str, config: Dict[str, Any]) -> Series[Any]:
         endpoint = f"{self._endpoint_prefix()}{predict_mode}.estimate"
@@ -79,7 +76,7 @@ def type(self) -> str:
             The type of the model.
 
         """
-        return self._list_info()["modelInfo"][0]["modelType"]  # type: ignore
+        return self._list_info()["modelInfo"]["modelType"]  # type: ignore
 
     def train_config(self) -> Series[Any]:
         """
@@ -89,7 +86,7 @@ def train_config(self) -> Series[Any]:
             The train config of the model.
 
         """
-        train_config: Series[Any] = Series(self._list_info()["trainConfig"][0])
+        train_config: Series[Any] = Series(self._list_info()["trainConfig"])
         return train_config
 
     def graph_schema(self) -> Series[Any]:
@@ -100,7 +97,7 @@ def graph_schema(self) -> Series[Any]:
             The graph schema of the model.
 
         """
-        graph_schema: Series[Any] = Series(self._list_info()["graphSchema"][0])
+        graph_schema: Series[Any] = Series(self._list_info()["graphSchema"])
         return graph_schema
 
     def loaded(self) -> bool:
@@ -111,7 +108,7 @@ def loaded(self) -> bool:
             True if the model is loaded in memory, False otherwise.
 
         """
-        return self._list_info()["loaded"].squeeze()  # type: ignore
+        return self._list_info()["loaded"]  # type: ignore
 
     def stored(self) -> bool:
         """
@@ -121,7 +118,7 @@ def stored(self) -> bool:
             True if the model is stored on disk, False otherwise.
 
         """
-        return self._list_info()["stored"].squeeze()  # type: ignore
+        return self._list_info()["stored"]  # type: ignore
 
     def creation_time(self) -> Any:  # neo4j.time.DateTime not exported
         """
@@ -131,7 +128,7 @@ def creation_time(self) -> Any:  # neo4j.time.DateTime not exported
             The creation time of the model.
 
         """
-        return self._list_info()["creationTime"].squeeze()
+        return self._list_info()["creationTime"]
 
     def shared(self) -> bool:
         """
@@ -141,7 +138,7 @@ def shared(self) -> bool:
             True if the model is shared, False otherwise.
 
         """
-        return self._list_info()["shared"].squeeze()  # type: ignore
+        return self._list_info()["shared"]  # type: ignore
 
     @compatible_with("published", min_inclusive=ServerVersion(2, 5, 0))
     def published(self) -> bool:
@@ -152,17 +149,17 @@ def published(self) -> bool:
             True if the model is published, False otherwise.
 
         """
-        return self._list_info()["published"].squeeze()  # type: ignore
+        return self._list_info()["published"]  # type: ignore
 
-    def model_info(self) -> Series[Any]:
+    def model_info(self) -> Dict[str, Any]:
         """
         Get the model info of the model.
 
         Returns:
             The model info of the model.
 
         """
-        return Series(self._list_info()["modelInfo"].squeeze())
+        return Series(self._list_info()["modelInfo"])  # type: ignore
 
     def exists(self) -> bool:
         """
@@ -199,11 +196,11 @@ def drop(self, failIfMissing: bool = False) -> Series[Any]:
                 "gds.beta.model.drop", params=params, custom_error=False
             ).squeeze()
         else:
-            result: Optional[Series[Any]] = self._query_runner.call_procedure(
+            result: Series[Any] = self._query_runner.call_procedure(
                 "gds.model.drop", params=params, custom_error=False
             ).squeeze()
 
-            if result is None:
+            if result.empty:
                 return Series()
 
             #  modelInfo {.*, modelName: modelName, modelType: modelType} AS modelInfo
@@ -223,7 +220,7 @@ def metrics(self) -> Series[Any]:
             The metrics of the model.
 
         """
-        model_info = self._list_info()["modelInfo"][0]
+        model_info = self._list_info()["modelInfo"]
         metrics: Series[Any] = Series(model_info["metrics"])
         return metrics
 

graphdatascience/model/node_classification_model.py

@@ -62,7 +62,7 @@ def classes(self) -> List[int]:
             The classes of the model.
 
         """
-        return self._list_info()["modelInfo"][0]["classes"]  # type: ignore
+        return self._list_info()["modelInfo"]["classes"]  # type: ignore
 
     def feature_properties(self) -> List[str]:
         """
@@ -72,5 +72,5 @@ def feature_properties(self) -> List[str]:
             The feature properties of the model.
 
         """
-        features: List[Dict[str, Any]] = self._list_info()["modelInfo"][0]["pipeline"]["featureProperties"]
+        features: List[Dict[str, Any]] = self._list_info()["modelInfo"]["pipeline"]["featureProperties"]
         return [f["feature"] for f in features]

graphdatascience/model/node_regression_model.py

@@ -21,5 +21,5 @@ def feature_properties(self) -> List[str]:
             The feature properties of the model.
 
         """
-        features: List[Dict[str, Any]] = self._list_info()["modelInfo"][0]["pipeline"]["featureProperties"]
+        features: List[Dict[str, Any]] = self._list_info()["modelInfo"]["pipeline"]["featureProperties"]
         return [f["feature"] for f in features]

graphdatascience/model/pipeline_model.py

@@ -75,7 +75,7 @@ def best_parameters(self) -> "Series[Any]":
             The best parameters for the pipeline model.
 
         """
-        best_params: Dict[str, Any] = self._list_info()["modelInfo"][0]["bestParameters"]
+        best_params: Dict[str, Any] = self._list_info()["modelInfo"]["bestParameters"]
         return Series(best_params)
 
     def node_property_steps(self) -> List[NodePropertyStep]:
@@ -86,7 +86,7 @@ def node_property_steps(self) -> List[NodePropertyStep]:
             The node property steps for the pipeline model.
 
         """
-        steps: List[Dict[str, Any]] = self._list_info()["modelInfo"][0]["pipeline"]["nodePropertySteps"]
+        steps: List[Dict[str, Any]] = self._list_info()["modelInfo"]["pipeline"]["nodePropertySteps"]
         return [NodePropertyStep(s["name"], s["config"]) for s in steps]
 
     def metrics(self) -> "Series[Any]":
@@ -97,6 +97,6 @@ def metrics(self) -> "Series[Any]":
             The metrics for the pipeline model.
 
         """
-        model_metrics: Dict[str, Any] = self._list_info()["modelInfo"][0]["metrics"]
+        model_metrics: Dict[str, Any] = self._list_info()["modelInfo"]["metrics"]
         metric_scores: Dict[str, MetricScores] = {k: MetricScores.create(v) for k, v in (model_metrics.items())}
         return Series(metric_scores)

graphdatascience/tests/integration/test_model_ops.py

@@ -305,6 +305,18 @@ def test_model_get_no_model(gds: GraphDataScience) -> None:
         gds.model.get("no_model")
 
 
+def test_missing_model_info(gds: GraphDataScience) -> None:
+    model = GraphSageModel("ghost-model", gds._query_runner, gds.server_version())
+    with pytest.raises(ValueError, match="There is no 'ghost-model' in the model catalog"):
+        model.model_info()
+
+
+def test_missing_model_drop(gds: GraphDataScience) -> None:
+    model = GraphSageModel("ghost-model", gds._query_runner, gds.server_version())
+
+    assert model.drop(failIfMissing=False).empty
+
+
 @pytest.mark.model_store_location
 @pytest.mark.filterwarnings("ignore: The query used a deprecated procedure.")
 def test_deprecated_model_Functions_still_work(gds: GraphDataScience, gs_model: GraphSageModel) -> None: