Changes so that the run method is not required anymore

Author: stellasia

src/neo4j_graphrag/experimental/components/entity_relation_extractor.py

@@ -14,7 +14,6 @@
 #  limitations under the License.
 from __future__ import annotations
 
-import abc
 import asyncio
 import enum
 import json
@@ -115,7 +114,7 @@ def fix_invalid_json(raw_json: str) -> str:
     return repaired_json
 
 
-class EntityRelationExtractor(Component, abc.ABC):
+class EntityRelationExtractor(Component):
     """Abstract class for entity relation extraction components.
 
     Args:
@@ -133,15 +132,14 @@ def __init__(
         self.on_error = on_error
         self.create_lexical_graph = create_lexical_graph
 
-    @abc.abstractmethod
     async def run(
         self,
         chunks: TextChunks,
         document_info: Optional[DocumentInfo] = None,
         lexical_graph_config: Optional[LexicalGraphConfig] = None,
         **kwargs: Any,
     ) -> Neo4jGraph:
-        pass
+        raise NotImplementedError()
 
     def update_ids(
         self,

src/neo4j_graphrag/experimental/components/resolver.py

@@ -12,7 +12,6 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-import abc
 from typing import Any, Optional
 
 import neo4j
@@ -22,7 +21,7 @@
 from neo4j_graphrag.utils import driver_config
 
 
-class EntityResolver(Component, abc.ABC):
+class EntityResolver(Component):
     """Entity resolution base class
 
     Args:
@@ -38,9 +37,8 @@ def __init__(
         self.driver = driver_config.override_user_agent(driver)
         self.filter_query = filter_query
 
-    @abc.abstractmethod
     async def run(self, *args: Any, **kwargs: Any) -> ResolutionStats:
-        pass
+        raise NotImplementedError()
 
 
 class SinglePropertyExactMatchResolver(EntityResolver):

src/neo4j_graphrag/experimental/pipeline/component.py

@@ -14,7 +14,6 @@
 #  limitations under the License.
 from __future__ import annotations
 
-import abc
 import inspect
 from typing import Any, get_type_hints
 
@@ -31,18 +30,14 @@ class DataModel(BaseModel):
     pass
 
 
-class ComponentMeta(abc.ABCMeta):
+class ComponentMeta(type):
     def __new__(
         meta, name: str, bases: tuple[type, ...], attrs: dict[str, Any]
     ) -> type:
         # extract required inputs and outputs from the run method signature
         run_method = attrs.get("run")
         run_context_method = attrs.get("run_with_context")
         run = run_context_method or run_method
-        if run is None:
-            raise RuntimeError(
-                f"Either 'run' or 'run_with_context' must be implemented in component: '{name}'"
-            )
         sig = inspect.signature(run)
         attrs["component_inputs"] = {
             param.name: {
@@ -73,7 +68,7 @@ def __new__(
         return type.__new__(meta, name, bases, attrs)
 
 
-class Component(abc.ABC, metaclass=ComponentMeta):
+class Component(metaclass=ComponentMeta):
     """Interface that needs to be implemented
     by all components.
     """
@@ -84,12 +79,27 @@ class Component(abc.ABC, metaclass=ComponentMeta):
     component_inputs: dict[str, dict[str, str | bool]]
     component_outputs: dict[str, dict[str, str | bool | type]]
 
-    @abc.abstractmethod
     async def run(self, *args: Any, **kwargs: Any) -> DataModel:
-        pass
+        """This function is planned for deprecation in a future release.
+
+        Note: if `run_with_context` is implemented, this method will not be used.
+        """
+        raise NotImplementedError(
+            "You must implement the `run` or `run_with_context` method. "
+            "`run` method will be marked for deprecation in a future release."
+        )
 
     async def run_with_context(
         self, context_: RunContext, *args: Any, **kwargs: Any
     ) -> DataModel:
+        """This method is called by the pipeline orchestrator.
+        The `context_` parameter contains information about
+        the pipeline run: the `run_id` and a `notify` function
+        that can be used to send events from the component to
+        the pipeline callback.
+
+        For now, it defaults to calling the `run` method, but it
+        is meant to replace the `run` method in a future release.
+        """
         # default behavior to prevent a breaking change
         return await self.run(*args, **kwargs)

tests/unit/experimental/pipeline/components.py

@@ -45,9 +45,6 @@ async def run(self, number1: int, number2: int = 2) -> IntResultModel:
 
 
 class ComponentMultiplyWithContext(Component):
-    async def run(self, number1: int, number2: int) -> IntResultModel:
-        return IntResultModel(result=number1 * number2)
-
     async def run_with_context(
         self, context_: RunContext, number1: int, number2: int = 2
     ) -> IntResultModel: