Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 18 additions & 10 deletions python/flink_agents/plan/actions/action.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@
from flink_agents.plan.function import Function, JavaFunction, PythonFunction

_CONFIG_TYPE = "__config_type__"
# Tags a config entry that we serialized from a pydantic model, so on the way
# back we only rebuild the ones we tagged and leave plain values alone.
_PYDANTIC_MODEL_MARKER = "__pydantic_model__"


class Action(BaseModel):
Expand Down Expand Up @@ -70,11 +73,12 @@ def __serialize_config(self, config: Dict[str, Any]) -> Dict[str, Any] | None:
data[_CONFIG_TYPE] = "python"
for name, value in config.items():
if isinstance(value, BaseModel):
data[name] = (
inspect.getmodule(value).__name__,
value.__class__.__name__,
value,
)
data[name] = {
_PYDANTIC_MODEL_MARKER: True,
"module": inspect.getmodule(value).__name__,
"class": value.__class__.__name__,
"value": value,
}
else:
data[name] = value
return data
Expand All @@ -95,11 +99,15 @@ def __custom_deserialize(self) -> "Action":
self["config"][name] = value["value"]
return self
for name, value in config.items():
try:
module = importlib.import_module(value[0])
clazz = getattr(module, value[1])
self["config"][name] = clazz.model_validate(value[2])
except Exception: # noqa : PERF203
# Rebuild only entries with our marker, so a plain list or dict
# value is not treated as a model by mistake. Do not catch errors:
# if the class cannot be imported (e.g. missing on the worker),
# fail here with the real cause instead of a confusing error later.
if isinstance(value, dict) and value.get(_PYDANTIC_MODEL_MARKER):
module = importlib.import_module(value["module"])
clazz = getattr(module, value["class"])
self["config"][name] = clazz.model_validate(value["value"])
else:
self["config"][name] = value
return self

Expand Down
11 changes: 6 additions & 5 deletions python/flink_agents/plan/tests/resources/action.json
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,11 @@
],
"config": {
"__config_type__": "python",
"output_schema": [
"flink_agents.api.agents.types",
"OutputSchema",
{
"output_schema": {
"__pydantic_model__": true,
"module": "flink_agents.api.agents.types",
"class": "OutputSchema",
"value": {
"output_schema": {
"names": [
"result"
Expand All @@ -23,6 +24,6 @@
]
}
}
]
}
}
}
51 changes: 51 additions & 0 deletions python/flink_agents/plan/tests/test_action.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,3 +120,54 @@ def test_action_deserialize_java_shape_config_unwraps_primitives() -> None:
"rate": 1.5,
"label": "fast",
}


def test_action_deserialize_python_config_propagates_reconstruction_error() -> None:
"""A tagged model whose module is missing must fail loudly.

The entry is marked as a serialized model, so importing a non-existent
module raises instead of being silently swallowed.
"""
json_str = json.dumps(
{
"name": "legal",
"exec": {
"func_type": "PythonFunction",
"module": "flink_agents.plan.tests.test_action",
"qualname": "legal_signature",
},
"trigger_conditions": ["_input_event"],
"config": {
"__config_type__": "python",
"broken": {
"__pydantic_model__": True,
"module": "nonexistent_module_xyz",
"class": "SomeClass",
"value": {},
},
},
}
)
with pytest.raises(ModuleNotFoundError):
Action.model_validate_json(json_str)


def test_action_deserialize_python_config_preserves_plain_list() -> None:
"""A plain user list must survive untouched, not be mistaken for a model."""
json_str = json.dumps(
{
"name": "legal",
"exec": {
"func_type": "PythonFunction",
"module": "flink_agents.plan.tests.test_action",
"qualname": "legal_signature",
},
"trigger_conditions": ["_input_event"],
"config": {
"__config_type__": "python",
"hosts": ["host-a", "host-b", "host-c"],
},
}
)
action = Action.model_validate_json(json_str)
assert action.config == {"hosts": ["host-a", "host-b", "host-c"]}
Loading