Merge pull request #13 from ProKil/feature-forget-datatype-decorator-warning

Node Intialization with Friendly Warning

Author: XuhuiZhou

docs/applications/cogym.md

@@ -2,6 +2,6 @@
 
 [Co-Gym](https://github.com/SALT-NLP/collaborative-gym) is a framework for enabling and evaluating **Human-Agent Collaboration** created by Yijia Shao, Vinay Samuel, Yucheng Jiang, John Yang, and Diyi Yang.
 
-Co-Gym uses aact as the communciation infrastructure between humans, agents, and environments to enable asynchronous action taking and observation receiving. 
+Co-Gym uses aact as the communciation infrastructure between humans, agents, and environments to enable asynchronous action taking and observation receiving.
 
-Please visit their GitHub repo for details. 
+Please visit their GitHub repo for details.

examples/exception_examples/example.toml

@@ -0,0 +1,13 @@
+redis_url = "redis://localhost:6379/0" # required
+extra_modules = ["exception_node"]
+
+[[nodes]]
+node_name = "exception_node"
+node_class = "exception_node"
+
+[nodes.node_args.print_channel_types]
+"tick/secs/1" = "tick"
+
+[[nodes]]
+node_name = "tick"
+node_class = "tick"

examples/exception_examples/exception_node.py

@@ -0,0 +1,12 @@
+from typing import AsyncIterator
+from aact.nodes import Node, NodeFactory
+from aact.messages import Text
+from aact import Message
+
+
+@NodeFactory.register("exception_node")
+class ExceptionNode(Node[Text, Text]):
+    def event_handler(
+        self, input_channel: str, input_message: Message[Text]
+    ) -> AsyncIterator[tuple[str, Message[Text]]]:
+        raise Exception("This is an exception from the node.")

src/aact/nodes/base.py

@@ -3,7 +3,7 @@
 import logging
 
 from ..utils import Self
-from typing import Any, AsyncIterator, Generic, Type, TypeVar
+from typing import Any, AsyncIterator, Generic, Literal, Type, TypeVar
 from pydantic import BaseModel, ConfigDict, ValidationError
 
 from abc import abstractmethod
@@ -176,6 +176,30 @@ def __init__(
         redis_url: str = "redis://localhost:6379/0",
     ):
         try:
+            for _, input_channel_type in input_channel_types:
+                if not issubclass(input_channel_type, DataModel):
+                    raise TypeError(
+                        f"Input channel type {input_channel_type} is not a subclass of DataModel"
+                    )
+                elif (
+                    input_channel_type.model_fields["data_type"].annotation  # type: ignore[comparison-overlap]
+                    == Literal[""]
+                ):
+                    raise TypeError(
+                        f"The input channel type {input_channel_type} needs to be registered with `@DataModelFactory.register`"
+                    )
+            for _, output_channel_type in output_channel_types:
+                if not issubclass(output_channel_type, DataModel):
+                    raise TypeError(
+                        f"Output channel type {output_channel_type} is not a subclass of DataModel"
+                    )
+                elif (
+                    output_channel_type.model_fields["data_type"].annotation  # type: ignore[comparison-overlap]
+                    == Literal[""]
+                ):
+                    raise TypeError(
+                        f"The output channel type {output_channel_type} needs to be registered with `@DataModelFactory.register`"
+                    )
             BaseModel.__init__(
                 self,
                 input_channel_types=dict(input_channel_types),

tests/messages/test_message.py

@@ -0,0 +1,19 @@
+from aact.messages import DataModel, DataModelFactory
+
+
+def test_create_data_model() -> None:
+    # Create a new data model
+    @DataModelFactory.register("MyDataModel")
+    class MyDataModel(DataModel):
+        name: str
+        age: int
+
+    # Create an instance of the data model
+    instance = MyDataModel(name="John", age=30)
+
+    # Validate the instance
+    assert instance.name == "John"
+    assert instance.age == 30
+
+    # Check if the instance is of type DataModel
+    assert isinstance(instance, DataModel)

tests/nodes/test_node_creation.py

@@ -0,0 +1,107 @@
+import pytest
+from aact.messages.base import DataModel
+from aact.messages.registry import DataModelFactory
+from aact.nodes import Node
+from aact.nodes import NodeFactory
+
+
+def test_node_with_non_datamodel_types() -> None:
+    """Test that creating a node with non-DataModel types raises a TypeError."""
+
+    class AWrong:
+        pass
+
+    class BWrong:
+        pass
+
+    # Create a node with non-DataModel types
+    @NodeFactory.register("MyNode")
+    class MyNode(Node[AWrong, BWrong]):  # type: ignore[type-var]
+        def __init__(self) -> None:
+            super().__init__(
+                node_name="MyNode",
+                input_channel_types=[("input", AWrong)],
+                output_channel_types=[("output", BWrong)],
+            )
+
+        async def event_handler(  # type: ignore[override]
+            self, input_channel: str, input_message: AWrong
+        ) -> None:
+            # Handle the event
+            pass
+
+    # Verify the correct error is raised
+    with pytest.raises(TypeError) as excinfo:
+        _ = MyNode()
+
+    assert (
+        "Input channel type <class 'test_node_creation.test_node_with_non_datamodel_types.<locals>.AWrong'> "
+        "is not a subclass of DataModel" in str(excinfo.value)
+    )
+
+
+def test_node_with_unregistered_datamodel_types() -> None:
+    """Test that creating a node with unregistered DataModel types raises a TypeError."""
+
+    class ACorrect(DataModel):
+        pass
+
+    class BCorrect(DataModel):
+        pass
+
+    # Create a node with unregistered DataModel types
+    @NodeFactory.register("MyNodeCorrect")
+    class MyNodeCorrect(Node[ACorrect, BCorrect]):
+        def __init__(self) -> None:
+            super().__init__(
+                node_name="MyNode",
+                input_channel_types=[("input", ACorrect)],
+                output_channel_types=[("output", BCorrect)],
+            )
+
+        async def event_handler(  # type: ignore[override]
+            self, input_channel: str, input_message: ACorrect
+        ) -> None:
+            # Handle the event
+            pass
+
+    # Verify the correct error is raised
+    with pytest.raises(TypeError) as excinfo:
+        _ = MyNodeCorrect()
+
+    assert (
+        "The input channel type <class 'test_node_creation.test_node_with_unregistered_datamodel_types.<locals>.ACorrect'> "
+        "needs to be registered with `@DataModelFactory.register`" in str(excinfo.value)
+    )
+
+
+def test_node_with_registered_datamodel_types() -> None:
+    """Test that creating a node with properly registered DataModel types succeeds."""
+
+    @DataModelFactory.register("ACorrectNew")
+    class ACorrectNew(DataModel):
+        pass
+
+    @DataModelFactory.register("BCorrectNew")
+    class BCorrectNew(DataModel):
+        pass
+
+    # Create a node with registered DataModel types
+    @NodeFactory.register("MyNodeCorrect")
+    class MyNodeCorrect(Node[ACorrectNew, BCorrectNew]):
+        def __init__(self) -> None:
+            super().__init__(
+                node_name="MyNode",
+                input_channel_types=[("input", ACorrectNew)],
+                output_channel_types=[("output", BCorrectNew)],
+            )
+
+        async def event_handler(  # type: ignore[override]
+            self, input_channel: str, input_message: ACorrectNew
+        ) -> None:
+            # Handle the event
+            pass
+
+    # Verify that node creation succeeds
+    node = MyNodeCorrect()
+    assert isinstance(node, MyNodeCorrect)