address comments.

Author: wenjin272

python/flink_agents/api/event.py

@@ -15,7 +15,6 @@
 #  See the License for the specific language governing permissions and
 # limitations under the License.
 #################################################################################
-import json
 from abc import ABC
 from typing import Any
 from uuid import UUID, uuid4
@@ -36,13 +35,15 @@ class Event(BaseModel, ABC, extra="allow"):
 
     @model_validator(mode='after')
     def validate_extra(self) -> 'Event':
-        """Make sure all extra properties are serializable."""
-        for value in self.__pydantic_extra__.values():
-            if isinstance(value, BaseModel):
-                continue
-            json.dumps(value)
+        """Ensure init fields is serializable."""
+        self.model_dump_json()
         return self
 
+    def __setattr__(self, name: str, value: Any) -> None:
+        super().__setattr__(name, value)
+        # Ensure added property can be serialized.
+        self.model_dump_json()
+
 
 class InputEvent(Event):
     """Event generated by the framework, carrying an input data that

python/flink_agents/api/tests/test_event.py

@@ -18,16 +18,25 @@
 from typing import Type
 
 import pytest
+from pydantic import ValidationError
+from pydantic_core import PydanticSerializationError
 
-from flink_agents.api.event import Event
+from flink_agents.api.event import Event, InputEvent, OutputEvent
 
 
-def test_event_serializable() -> None: #noqa D103
-    event = Event(a="1")
-    event.model_dump_json()
+def test_event_serializable_valid() -> None: #noqa D103
+    Event(a=1, b=InputEvent(input=1), c=OutputEvent(output='111'))
 
-    with pytest.raises(TypeError):
-        Event(a=Type[Event])
+def test_event_serializable_invalid() -> None: #noqa D103
+    with pytest.raises(ValidationError):
+        Event(a=1, b=Type[InputEvent])
+
+def test_event_add_field() -> None: #noqa D103
+    event = Event(a=1)
+    event.c = Event()
+
+def test_event_add_field_invalid() -> None: #noqa D103
+    event = Event(a=1)
+    with pytest.raises(PydanticSerializationError):
+        event.c = Type[InputEvent]
 
-    event = Event(a=Event())
-    event.model_dump_json()

python/flink_agents/plan/action.py

@@ -15,8 +15,7 @@
 #  See the License for the specific language governing permissions and
 # limitations under the License.
 #################################################################################
-import inspect
-from typing import Callable, List, Type
+from typing import List, Type
 
 from pydantic import BaseModel
 
@@ -41,6 +40,7 @@ class Action(BaseModel):
     """
 
     name: str
+    #TODO: Raise a warning when the action has a return value, as it will be ignored.
     exec: Function
     listen_event_types: List[Type[Event]]
 
@@ -52,19 +52,5 @@ def __init__(
     ) -> None:
         """Action will check function signature when init."""
         super().__init__(name=name, exec=exec, listen_event_types=listen_event_types)
-        exec.check_signature(self.check_signature)
-
-    @classmethod
-    def check_signature(cls, func: Callable) -> None:
-        """" Checker for action function signature."""
-        #TODO: update check logic after import State and RunnerContext.
-        params = inspect.signature(func).parameters
-        if len(params) != 1:
-            err_msg = "Action function must have exactly 1 parameter"
-            raise TypeError(err_msg)
-        for i, param in enumerate(params.values()):
-            if i == 0:
-                if not issubclass(param.annotation, Event):
-                    err_msg = "Action function first parameter must be Event"
-                    raise TypeError(err_msg)
+        exec.check_signature([Event])
 

python/flink_agents/plan/function.py

@@ -28,7 +28,7 @@ class Function(BaseModel, ABC):
     """Base interface for user defined functions, includes python and java."""
 
     @abstractmethod
-    def check_signature(self, checker: Callable) -> None:
+    def check_signature(self, *args: Tuple[Any, ...]) -> None:
         """Check function signature is legal or not."""
 
     @abstractmethod
@@ -76,9 +76,16 @@ def from_callable(func: Callable) -> Function:
             __func=func,
         )
 
-    def check_signature(self, checker: Callable) -> None:
-        """Apply external check logic to function signature."""
-        checker(self.__get_func())
+    def check_signature(self, *args: Tuple[Any, ...]) -> None:
+        """Check function signature."""
+        params = inspect.signature(self.__get_func()).parameters
+        annotations = [param.annotation for param in params.values()]
+        err_msg = f"Expect {self.qualname} have signature {args}, but got {annotations}."
+        if len(params) != args.__len__():
+            raise TypeError(err_msg)
+        for i, annotation in enumerate(annotations):
+            if not issubclass(annotation, *args[i]):
+                raise TypeError(err_msg)
 
     def __call__(self, *args: Tuple[Any, ...], **kwargs: Dict[str, Any]) -> Any:
         """Execute the stored function with provided arguments.
@@ -119,11 +126,12 @@ def __get_func(self) -> Callable:
         return self.__func
 
 
+#TODO: Implement JavaFunction.
 class JavaFunction(Function):
     """Descriptor for a java callable function."""
 
     def __call__(self, *args: Tuple[Any, ...], **kwargs: Dict[str, Any]) -> Any:
         """Execute the stored function with provided arguments."""
 
-    def check_signature(self, checker: Callable) -> None:
+    def check_signature(self, *args: Tuple[Any, ...]) -> None:
         """Check function signature is legal or not."""

python/flink_agents/plan/tests/test_action.py

@@ -17,31 +17,30 @@
 #################################################################################
 import pytest
 
-from flink_agents.api.event import InputEvent, OutputEvent
+from flink_agents.api.event import InputEvent
 from flink_agents.plan.action import Action
 from flink_agents.plan.function import PythonFunction
 
 
-def increment(event: InputEvent) -> OutputEvent: # noqa: D103
-    value = event.input
-    value += 1
-    return OutputEvent(output=value)
+def legal_signature(event: InputEvent) -> None: # noqa: D103
+    pass
 
-def decrement(value: int) -> OutputEvent: # noqa: D103
-    value -= 1
-    return OutputEvent(output=value)
+def illegal_signature(value: int) ->  None: # noqa: D103
+    pass
 
-def test_action_signature() -> None: # noqa: D103
+def test_action_signature_legal() -> None: # noqa: D103
     Action(
-        name="increment",
-        exec=PythonFunction.from_callable(increment),
+        name="legal",
+        exec=PythonFunction.from_callable(legal_signature),
         listen_event_types=[InputEvent],
     )
 
+def test_action_signature_illegal() -> None:  # noqa: D103
     with pytest.raises(TypeError):
         Action(
-            name="decrement",
-            exec=PythonFunction.from_callable(decrement),
+            name="illegal",
+            exec=PythonFunction.from_callable(illegal_signature),
             listen_event_types=[InputEvent],
         )
 
+