Support GraphSage train + Model catalog ops in v2 endpoints

Author: FlorentinD

graphdatascience/model/v2/graphsage_model.py

@@ -1,11 +1,12 @@
 from typing import Any
 
 from pandas import Series
+from pydantic import BaseModel
+from pydantic.alias_generators import to_camel
 
-from ...call_parameters import CallParameters
 from ...graph.graph_object import Graph
 from ...graph.graph_type_check import graph_type_check
-from ..model import Model
+from .model import Model
 
 
 class GraphSageModelV2(Model):
@@ -30,13 +31,7 @@ def predict_write(self, G: Graph, **config: Any) -> "Series[Any]":
             The result of the write operation.
 
         """
-        endpoint = self._endpoint_prefix() + "write"
-        config["modelName"] = self.name()
-        params = CallParameters(graph_name=G.name(), config=config)
-
-        return self._query_runner.call_procedure(  # type: ignore
-            endpoint=endpoint, params=params, logging=True
-        ).squeeze()
+        raise ValueError
 
     @graph_type_check
     def predict_write_estimate(self, G: Graph, **config: Any) -> "Series[Any]":
@@ -51,4 +46,28 @@ def predict_write_estimate(self, G: Graph, **config: Any) -> "Series[Any]":
             The memory needed to generate embeddings for the given graph and write the results to the database.
 
         """
-        return self._estimate_predict("write", G.name(), config)
+        raise ValueError
+
+
+class GraphSageMutateResult(BaseModel, alias_generator=to_camel):
+    node_count: int
+    node_properties_written: int
+    pre_processing_millis: int
+    compute_millis: int
+    mutate_millis: int
+    configuration: dict[str, Any]
+
+    def __getitem__(self, item: str) -> Any:
+        return self.__dict__[item]
+
+
+class GraphSageWriteResult(BaseModel, alias_generator=to_camel):
+    node_count: int
+    node_properties_written: int
+    pre_processing_millis: int
+    compute_millis: int
+    write_millis: int
+    configuration: dict[str, Any]
+
+    def __getitem__(self, item: str) -> Any:
+        return self.__dict__[item]

graphdatascience/model/v2/model.py

@@ -1,32 +1,17 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
-from datetime import datetime
-from typing import Any
+from abc import ABC
+from typing import Optional
 
-from pandas import DataFrame, Series
-
-from graphdatascience.model.v2.model_info import ModelInfo
-
-from ..call_parameters import CallParameters
-from ..graph.graph_object import Graph
-from ..graph.graph_type_check import graph_type_check
-from ..query_runner.query_runner import QueryRunner
-from ..server_version.compatible_with import compatible_with
-from ..server_version.server_version import ServerVersion
-
-
-class InfoProvider(ABC):
-    @abstractmethod
-    def fetch(self, model_name: str) -> ModelInfo:
-        """Return the task with progress for the given job_id."""
-        pass
+from graphdatascience.model.v2.model_api import ModelApi
+from graphdatascience.model.v2.model_info import ModelDetails
 
 
+# Compared to v1 Model offering typed parameters for predict endpoints
 class Model(ABC):
-    def __init__(self, name: str, info_provider: InfoProvider):
+    def __init__(self, name: str, model_api: ModelApi):
         self._name = name
-        self._info_provider = info_provider
+        self._model_api = model_api
 
     # TODO estimate mode, predict modes on here?
     # implement Cypher and Arrow info_provider and stuff
@@ -41,95 +26,8 @@ def name(self) -> str:
         """
         return self._name
 
-    def type(self) -> str:
-        """
-        Get the type of the model.
-
-        Returns:
-            The type of the model.
-
-        """
-        return self._info_provider.fetch(self._name).type
-
-    def train_config(self) -> Series[Any]:
-        """
-        Get the train config of the model.
-
-        Returns:
-            The train config of the model.
-
-        """
-        return self._info_provider.fetch(self._name).train_config
-
-    def graph_schema(self) -> Series[Any]:
-        """
-        Get the graph schema of the model.
-
-        Returns:
-            The graph schema of the model.
-
-        """
-        return self._info_provider.fetch(self._name).graph_schema
-
-    def loaded(self) -> bool:
-        """
-        Check whether the model is loaded in memory.
-
-        Returns:
-            True if the model is loaded in memory, False otherwise.
-
-        """
-        return self._info_provider.fetch(self._name).loaded
-
-    def stored(self) -> bool:
-        """
-        Check whether the model is stored on disk.
-
-        Returns:
-            True if the model is stored on disk, False otherwise.
-
-        """
-        return self._info_provider.fetch(self._name).stored
-
-    def creation_time(self) -> datetime.datetime:
-        """
-        Get the creation time of the model.
-
-        Returns:
-            The creation time of the model.
-
-        """
-        return self._info_provider.fetch(self._name).creation_time
-
-    def shared(self) -> bool:
-        """
-        Check whether the model is shared.
-
-        Returns:
-            True if the model is shared, False otherwise.
-
-        """
-        return self._info_provider.fetch(self._name).shared
-
-    def published(self) -> bool:
-        """
-        Check whether the model is published.
-
-        Returns:
-            True if the model is published, False otherwise.
-
-        """
-        return self._info_provider.fetch(self._name).published
-
-    def model_info(self) -> dict[str, Any]:
-        """
-        Get the model info of the model.
-
-        Returns:
-            The model info of the model.
-
-        """
-        return self._info_provider.fetch(self._name).model_info
+    def details(self) -> ModelDetails:
+        return self._model_api.get(self._name)
 
     def exists(self) -> bool:
         """
@@ -139,9 +37,9 @@ def exists(self) -> bool:
             True if the model exists, False otherwise.
 
         """
-        raise NotImplementedError()
+        return self._model_api.exists(self._name)
 
-    def drop(self, failIfMissing: bool = False) -> Series[Any]:
+    def drop(self, failIfMissing: bool = False) -> Optional[ModelDetails]:
         """
         Drop the model.
 
@@ -152,22 +50,10 @@ def drop(self, failIfMissing: bool = False) -> Series[Any]:
             The result of the drop operation.
 
         """
-        raise NotImplementedError()
-
-    def metrics(self) -> Series[Any]:
-        """
-        Get the metrics of the model.
-
-        Returns:
-            The metrics of the model.
-
-        """
-        model_info = self._info_provider.fetch(self._name).model_info
-        metrics: Series[Any] = Series(model_info["metrics"])
-        return metrics
+        return self._model_api.drop(self._name, failIfMissing)
 
     def __str__(self) -> str:
-        return f"{self.__class__.__name__}(name={self.name()}, type={self.type()})"
+        return f"{self.__class__.__name__}(name={self.name()}, type={self.details().type})"
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self._info_provider.fetch(self._name).to_dict()})"
+        return f"{self.__class__.__name__}({self.details().model_dump()})"

graphdatascience/model/v2/model_api.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from graphdatascience.model.v2.model_info import ModelDetails
+
+
+class ModelApi(ABC):
+    """
+    Abstract base class defining the API for model operations.
+    This class is intended to be subclassed by specific model implementations.
+    """
+
+    @abstractmethod
+    def exists(self, model: str) -> bool:
+        """
+        Check if a specific model exists.
+
+        Args:
+            model: The name of the model.
+
+        Returns:
+            True if the model exists, False otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def get(self, model: str) -> ModelDetails:
+        """
+        Get the details of a specific model.
+
+        Args:
+            model: The name of the model.
+
+        Returns:
+            The details of the model.
+        """
+        pass
+
+    @abstractmethod
+    def drop(self, model: str, fail_if_missing: bool) -> Optional[ModelDetails]:
+        """
+        Drop a specific model.
+
+        Args:
+            model: The name of the model.
+            fail_if_missing: If True, an error is thrown if the model does not exist. If False, no error is thrown.
+
+        Returns:
+            The result of the drop operation.
+        """
+        pass

graphdatascience/model/v2/model_info.py

@@ -1,21 +1,20 @@
 import datetime
 from typing import Any
-from pydantic import BaseModel
-from abc import ABC, abstractmethod
+
+from pydantic import BaseModel, Field
 from pydantic.alias_generators import to_camel
 
 
-class ModelInfo(BaseModel, alias_generator=to_camel):
-    name: str
-    type: str
+class ModelDetails(BaseModel, alias_generator=to_camel):
+    name: str = Field(alias="modelName")
+    type: str = Field(alias="modelType")
     train_config: dict[str, Any]
     graph_schema: dict[str, Any]
     loaded: bool
     stored: bool
-    shared: bool
     published: bool
     model_info: dict[str, Any]  # TODO better typing in actual model?
-    creation_time: datetime.datetime  # TODO correct type? / conversion needed
+    creation_time: datetime.datetime
 
     def __getitem__(self, item: str) -> Any:
         return getattr(self, item)