282 lines
8.4 KiB
Python
282 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from typing import Annotated, Literal, TypeAliasType
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
|
|
|
from wf_core.models.json_values import JsonValue, validate_strict_json_value
|
|
from wf_core.paths import GraphSourcePath, LocalPath
|
|
|
|
MAX_INPUT_EXPRESSION_DEPTH = 64
|
|
MAX_INPUT_EXPRESSION_NODES = 1_024
|
|
|
|
|
|
class InputPathBinding(BaseModel):
|
|
"""Map one workflow graph source path into one node-local input path."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
target: LocalPath = Field(
|
|
description=(
|
|
"Node-local input path to populate. Prefer canonical strings such "
|
|
"as `field` or `.` for the whole node input payload. Structural "
|
|
"objects such as {'root': 'local', 'parts': ['field']} are also "
|
|
"accepted as input."
|
|
)
|
|
)
|
|
path: GraphSourcePath = Field(
|
|
description=(
|
|
"Workflow source path to read from input, state, or context. "
|
|
"Prefer canonical strings such as `input.text` or `state.report`. "
|
|
"Structural objects such as {'root': 'input', 'parts': ['text']} "
|
|
"are also accepted as input."
|
|
)
|
|
)
|
|
|
|
|
|
class InputValueBinding(BaseModel):
|
|
"""Map one static value into one node-local input path."""
|
|
|
|
model_config = ConfigDict(extra="forbid", strict=True)
|
|
|
|
target: LocalPath = Field(
|
|
description="Node-local input path that receives this literal JSON value."
|
|
)
|
|
value: JsonValue = Field(
|
|
description=(
|
|
"Literal JSON-compatible value to pass to the node. Use this for "
|
|
"constants, not for values read from workflow input or state."
|
|
)
|
|
)
|
|
|
|
@field_validator("value", mode="before")
|
|
@classmethod
|
|
def validate_value(cls, value: object) -> JsonValue:
|
|
return validate_strict_json_value(value)
|
|
|
|
|
|
class LiteralExpression(BaseModel):
|
|
"""A strict JSON literal embedded in a node-local input expression."""
|
|
|
|
model_config = ConfigDict(extra="forbid", strict=True)
|
|
|
|
kind: Literal["literal"]
|
|
value: JsonValue
|
|
|
|
@field_validator("value", mode="before")
|
|
@classmethod
|
|
def validate_value(cls, value: object) -> JsonValue:
|
|
return validate_strict_json_value(value)
|
|
|
|
|
|
class PathExpression(BaseModel):
|
|
"""Read one graph source path while resolving a composite input."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
kind: Literal["path"]
|
|
path: GraphSourcePath
|
|
|
|
|
|
class ArrayExpression(BaseModel):
|
|
"""Resolve ordered child expressions into one JSON array."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
kind: Literal["array"]
|
|
items: list[InputExpression]
|
|
|
|
|
|
class ObjectExpression(BaseModel):
|
|
"""Resolve named child expressions into one JSON object."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
kind: Literal["object"]
|
|
fields: dict[str, InputExpression]
|
|
|
|
|
|
type InputExpression = Annotated[
|
|
LiteralExpression | PathExpression | ArrayExpression | ObjectExpression,
|
|
Field(discriminator="kind"),
|
|
]
|
|
|
|
|
|
class OccurrenceExpression(BaseModel):
|
|
"""Reference one typed schedule-occurrence field.
|
|
|
|
Schedule-only leaf: graph expressions must not accept this kind, and
|
|
``GraphSourcePath`` roots stay closed to input/state/context.
|
|
"""
|
|
|
|
model_config = ConfigDict(extra="forbid", strict=True)
|
|
|
|
kind: Literal["occurrence"]
|
|
field: Literal["schedule_id", "occurrence_id", "scheduled_at"]
|
|
|
|
|
|
class ScheduleArrayExpression(BaseModel):
|
|
"""Resolve ordered child schedule expressions into one JSON array."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
kind: Literal["array"]
|
|
items: list[ScheduleExpression]
|
|
|
|
|
|
class ScheduleObjectExpression(BaseModel):
|
|
"""Resolve named child schedule expressions into one JSON object."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
kind: Literal["object"]
|
|
fields: dict[str, ScheduleExpression]
|
|
|
|
|
|
type ScheduleExpression = Annotated[
|
|
LiteralExpression
|
|
| OccurrenceExpression
|
|
| ScheduleArrayExpression
|
|
| ScheduleObjectExpression,
|
|
Field(discriminator="kind"),
|
|
]
|
|
|
|
|
|
def _raise_limit(limit: str, location: str) -> None:
|
|
raise ValueError(f"input expression {limit} limit exceeded at {location}")
|
|
|
|
|
|
def validate_input_expression_limits(
|
|
value: object,
|
|
*,
|
|
max_depth: int = MAX_INPUT_EXPRESSION_DEPTH,
|
|
max_nodes: int = MAX_INPUT_EXPRESSION_NODES,
|
|
) -> None:
|
|
"""Bound raw expression trees before Pydantic recursively constructs them."""
|
|
|
|
nodes = 0
|
|
|
|
def visit_json(raw: object, *, depth: int, location: str) -> None:
|
|
nonlocal nodes
|
|
if isinstance(raw, list):
|
|
visit_container(depth=depth, location=location)
|
|
for index, item in enumerate(raw):
|
|
visit_json(item, depth=depth + 1, location=f"{location}[{index}]")
|
|
elif isinstance(raw, Mapping):
|
|
visit_container(depth=depth, location=location)
|
|
for key, item in raw.items():
|
|
if isinstance(key, str):
|
|
visit_json(item, depth=depth + 1, location=f"{location}.{key}")
|
|
|
|
def visit_container(*, depth: int, location: str) -> None:
|
|
nonlocal nodes
|
|
if depth > max_depth:
|
|
_raise_limit("depth", location)
|
|
nodes += 1
|
|
if nodes > max_nodes:
|
|
_raise_limit("node", location)
|
|
|
|
def visit_expression(raw: object, *, depth: int, location: str) -> None:
|
|
nonlocal nodes
|
|
if depth > max_depth:
|
|
_raise_limit("depth", location)
|
|
nodes += 1
|
|
if nodes > max_nodes:
|
|
_raise_limit("node", location)
|
|
if not isinstance(raw, Mapping):
|
|
return
|
|
|
|
kind = raw.get("kind")
|
|
if kind == "literal":
|
|
visit_json(raw.get("value"), depth=depth + 1, location=f"{location}.value")
|
|
elif kind == "array" and isinstance(raw.get("items"), list):
|
|
for index, item in enumerate(raw["items"]):
|
|
visit_expression(
|
|
item, depth=depth + 1, location=f"{location}.items[{index}]"
|
|
)
|
|
elif kind == "object" and isinstance(raw.get("fields"), Mapping):
|
|
for key, item in raw["fields"].items():
|
|
if isinstance(key, str):
|
|
visit_expression(
|
|
item, depth=depth + 1, location=f"{location}.fields.{key}"
|
|
)
|
|
|
|
visit_expression(value, depth=1, location="expression")
|
|
|
|
|
|
def _check_expression_limits(value: object) -> object:
|
|
"""Share the budget validator across graph and schedule bindings."""
|
|
raw_value = (
|
|
value.model_dump(mode="python") if isinstance(value, BaseModel) else value
|
|
)
|
|
validate_input_expression_limits(raw_value)
|
|
return value
|
|
|
|
|
|
class InputExpressionBinding(BaseModel):
|
|
"""Assign one recursively composed expression to a node-local target."""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
target: LocalPath
|
|
expression: InputExpression
|
|
|
|
@field_validator("expression", mode="before")
|
|
@classmethod
|
|
def check_limits(cls, value: object) -> object:
|
|
return _check_expression_limits(value)
|
|
|
|
|
|
InputBinding = Annotated[
|
|
InputPathBinding | InputValueBinding,
|
|
Field(
|
|
union_mode="left_to_right",
|
|
description=(
|
|
"Simple canonical binding for node inputs or workflow outputs. Use "
|
|
"either a path binding with `path`, or a literal binding with `value`; "
|
|
"composite `expression` bindings are node-local only."
|
|
),
|
|
),
|
|
]
|
|
"""Canonical simple node input binding, distinguished by `path` vs `value`."""
|
|
|
|
|
|
StepInputBinding = TypeAliasType(
|
|
"StepInputBinding",
|
|
Annotated[
|
|
InputPathBinding | InputValueBinding | InputExpressionBinding,
|
|
Field(union_mode="left_to_right"),
|
|
],
|
|
)
|
|
|
|
|
|
class ScheduleInputBinding(BaseModel):
|
|
"""Assign one schedule-side expression to a workflow-input target.
|
|
|
|
Schedule expressions reuse literal/object/array composition and the same
|
|
budget validator, but leaves are typed occurrence references only. Graph
|
|
paths are invalid here.
|
|
"""
|
|
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
target: LocalPath
|
|
expression: ScheduleExpression
|
|
|
|
@field_validator("expression", mode="before")
|
|
@classmethod
|
|
def check_limits(cls, value: object) -> object:
|
|
return _check_expression_limits(value)
|
|
|
|
|
|
for _model in (
|
|
ArrayExpression,
|
|
ObjectExpression,
|
|
InputExpressionBinding,
|
|
ScheduleArrayExpression,
|
|
ScheduleObjectExpression,
|
|
ScheduleInputBinding,
|
|
):
|
|
_model.model_rebuild()
|