in the folders 3
This commit is contained in:
@@ -11,8 +11,11 @@ from .callables import (
|
||||
SyncRegistryHandler,
|
||||
)
|
||||
from .inference import accepts_context, infer_models, is_basemodel_subclass
|
||||
from .decorator import node
|
||||
from .registry import build_async_registry, build_registry
|
||||
from .result import NodeReturn
|
||||
from .spec import NodeSpec, build_async_registry, build_registry, node
|
||||
from .schema import schema_ref_for
|
||||
from .spec import NodeSpec
|
||||
|
||||
__all__ = [
|
||||
"AsyncContextNodeCallable",
|
||||
@@ -33,4 +36,5 @@ __all__ = [
|
||||
"infer_models",
|
||||
"is_basemodel_subclass",
|
||||
"node",
|
||||
"schema_ref_for",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from inspect import iscoroutinefunction
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast, overload
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .callables import AsyncNodeCallable, InputT, NodeCallable, OutputT
|
||||
from .inference import accepts_context, infer_models
|
||||
from .spec import NodeSpec
|
||||
|
||||
|
||||
@overload
|
||||
def node(
|
||||
fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
|
||||
/,
|
||||
) -> NodeSpec[InputT, OutputT]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def node(
|
||||
fn: None = None,
|
||||
/,
|
||||
) -> Callable[
|
||||
[NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
|
||||
NodeSpec[InputT, OutputT],
|
||||
]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def node(
|
||||
fn: None = None,
|
||||
/,
|
||||
*,
|
||||
name: str | None = None,
|
||||
input_model: type[InputT] | None = None,
|
||||
output_model: type[OutputT] | None = None,
|
||||
outcomes: tuple[str, ...] = ("ok",),
|
||||
description: str | None = None,
|
||||
is_async: bool | None = None,
|
||||
) -> Callable[
|
||||
[NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
|
||||
NodeSpec[InputT, OutputT],
|
||||
]: ...
|
||||
|
||||
|
||||
def node(
|
||||
fn: NodeCallable[InputT, OutputT]
|
||||
| AsyncNodeCallable[InputT, OutputT]
|
||||
| None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
input_model: type[InputT] | None = None,
|
||||
output_model: type[OutputT] | None = None,
|
||||
outcomes: tuple[str, ...] = ("ok",),
|
||||
description: str | None = None,
|
||||
is_async: bool | None = None,
|
||||
) -> Any:
|
||||
"""Convert a typed Python function into a reusable workflow node spec."""
|
||||
def decorator(
|
||||
fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
|
||||
) -> NodeSpec[InputT, OutputT]:
|
||||
inferred_input_model: type[BaseModel] | None = input_model
|
||||
inferred_output_model: type[BaseModel] | None = output_model
|
||||
if inferred_input_model is None or inferred_output_model is None:
|
||||
inferred_input_model, inferred_output_model = infer_models(fn)
|
||||
|
||||
resolved_name = name or getattr(fn, "__name__", "node")
|
||||
resolved_is_async = iscoroutinefunction(fn) if is_async is None else is_async
|
||||
resolved_accepts_context = accepts_context(fn)
|
||||
return cast(
|
||||
NodeSpec[InputT, OutputT],
|
||||
NodeSpec(
|
||||
name=resolved_name,
|
||||
input_model=inferred_input_model,
|
||||
output_model=inferred_output_model,
|
||||
outcomes=outcomes,
|
||||
fn=cast(Any, fn),
|
||||
description=description or fn.__doc__,
|
||||
is_async=resolved_is_async,
|
||||
accepts_context=resolved_accepts_context,
|
||||
),
|
||||
)
|
||||
|
||||
if fn is not None:
|
||||
return decorator(fn)
|
||||
return decorator
|
||||
@@ -0,0 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal, overload
|
||||
|
||||
from .callables import AsyncRegistryHandler, SyncRegistryHandler
|
||||
from .spec import NodeSpec
|
||||
|
||||
|
||||
def build_registry(
|
||||
*specs: NodeSpec[Any, Any],
|
||||
) -> dict[str, SyncRegistryHandler]:
|
||||
"""Export node specs as sync runtime registry handlers."""
|
||||
return _build_registry(specs, export="sync")
|
||||
|
||||
|
||||
def build_async_registry(
|
||||
*specs: NodeSpec[Any, Any],
|
||||
) -> dict[str, AsyncRegistryHandler]:
|
||||
"""Export node specs as async runtime registry handlers."""
|
||||
return _build_registry(specs, export="async")
|
||||
|
||||
|
||||
@overload
|
||||
def _build_registry(
|
||||
specs: tuple[NodeSpec[Any, Any], ...],
|
||||
*,
|
||||
export: Literal["sync"],
|
||||
) -> dict[str, SyncRegistryHandler]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def _build_registry(
|
||||
specs: tuple[NodeSpec[Any, Any], ...],
|
||||
*,
|
||||
export: Literal["async"],
|
||||
) -> dict[str, AsyncRegistryHandler]: ...
|
||||
|
||||
|
||||
def _build_registry(
|
||||
specs: tuple[NodeSpec[Any, Any], ...],
|
||||
*,
|
||||
export: Literal["sync", "async"],
|
||||
) -> dict[str, Any]:
|
||||
if export == "sync":
|
||||
return {spec.name: spec.to_registry_handler() for spec in specs}
|
||||
if export == "async":
|
||||
return {spec.name: spec.to_async_registry_handler() for spec in specs}
|
||||
raise ValueError(f"unknown registry export mode {export!r}")
|
||||
@@ -10,5 +10,7 @@ OutputT_co = TypeVar("OutputT_co", bound=BaseModel, covariant=True)
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NodeReturn(Generic[OutputT_co]):
|
||||
"""A node result that explicitly selects the outgoing workflow outcome."""
|
||||
|
||||
outcome: str
|
||||
output: OutputT_co
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_core import SchemaRef
|
||||
|
||||
|
||||
def schema_ref_for(model_type: type[BaseModel]) -> SchemaRef:
|
||||
"""Build a core schema reference from a pydantic model class."""
|
||||
return SchemaRef.model_validate(model_type.model_json_schema())
|
||||
@@ -1,19 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from inspect import iscoroutinefunction
|
||||
from collections.abc import Awaitable, Callable
|
||||
from collections.abc import Awaitable
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
Any,
|
||||
Generic,
|
||||
Literal,
|
||||
cast,
|
||||
overload,
|
||||
)
|
||||
from typing import Any, Generic, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_core import NodeDef, RuntimeContext, SchemaRef
|
||||
from wf_core import NodeDef, RuntimeContext
|
||||
|
||||
from .callables import (
|
||||
AsyncNodeCallable,
|
||||
@@ -25,12 +18,8 @@ from .callables import (
|
||||
PlainNodeCallable,
|
||||
SyncRegistryHandler,
|
||||
)
|
||||
from .inference import accepts_context, infer_models
|
||||
from .result import NodeReturn
|
||||
|
||||
|
||||
def _schema_ref_for(model_type: type[BaseModel]) -> SchemaRef:
|
||||
return SchemaRef.model_validate(model_type.model_json_schema())
|
||||
from .schema import schema_ref_for
|
||||
|
||||
|
||||
def _default_outcome(spec: "NodeSpec[Any, Any]") -> str:
|
||||
@@ -61,6 +50,8 @@ def _coerce_registry_result(
|
||||
|
||||
@dataclass(slots=True)
|
||||
class NodeSpec(Generic[InputT, OutputT]):
|
||||
"""Authoring-time wrapper for a typed Python node function."""
|
||||
|
||||
name: str
|
||||
input_model: type[InputT]
|
||||
output_model: type[OutputT]
|
||||
@@ -84,8 +75,8 @@ class NodeSpec(Generic[InputT, OutputT]):
|
||||
def to_node_def(self) -> NodeDef:
|
||||
return NodeDef(
|
||||
name=self.name,
|
||||
input_schema=_schema_ref_for(self.input_model),
|
||||
output_schema=_schema_ref_for(self.output_model),
|
||||
input_schema=schema_ref_for(self.input_model),
|
||||
output_schema=schema_ref_for(self.output_model),
|
||||
outcomes=list(self.outcomes),
|
||||
)
|
||||
|
||||
@@ -129,119 +120,3 @@ class NodeSpec(Generic[InputT, OutputT]):
|
||||
)
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
@overload
|
||||
def node(
|
||||
fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
|
||||
/,
|
||||
) -> NodeSpec[InputT, OutputT]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def node(
|
||||
fn: None = None,
|
||||
/,
|
||||
) -> Callable[
|
||||
[NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
|
||||
NodeSpec[InputT, OutputT],
|
||||
]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def node(
|
||||
fn: None = None,
|
||||
/,
|
||||
*,
|
||||
name: str | None = None,
|
||||
input_model: type[InputT] | None = None,
|
||||
output_model: type[OutputT] | None = None,
|
||||
outcomes: tuple[str, ...] = ("ok",),
|
||||
description: str | None = None,
|
||||
is_async: bool | None = None,
|
||||
) -> Callable[
|
||||
[NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]],
|
||||
NodeSpec[InputT, OutputT],
|
||||
]: ...
|
||||
|
||||
|
||||
def node(
|
||||
fn: NodeCallable[InputT, OutputT]
|
||||
| AsyncNodeCallable[InputT, OutputT]
|
||||
| None = None,
|
||||
*,
|
||||
name: str | None = None,
|
||||
input_model: type[InputT] | None = None,
|
||||
output_model: type[OutputT] | None = None,
|
||||
outcomes: tuple[str, ...] = ("ok",),
|
||||
description: str | None = None,
|
||||
is_async: bool | None = None,
|
||||
) -> Any:
|
||||
def decorator(
|
||||
fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT],
|
||||
) -> NodeSpec[InputT, OutputT]:
|
||||
inferred_input_model: type[BaseModel] | None = input_model
|
||||
inferred_output_model: type[BaseModel] | None = output_model
|
||||
if inferred_input_model is None or inferred_output_model is None:
|
||||
inferred_input_model, inferred_output_model = infer_models(fn)
|
||||
|
||||
resolved_name = name or getattr(fn, "__name__", "node")
|
||||
resolved_is_async = iscoroutinefunction(fn) if is_async is None else is_async
|
||||
resolved_accepts_context = accepts_context(fn)
|
||||
return cast(
|
||||
NodeSpec[InputT, OutputT],
|
||||
NodeSpec(
|
||||
name=resolved_name,
|
||||
input_model=inferred_input_model,
|
||||
output_model=inferred_output_model,
|
||||
outcomes=outcomes,
|
||||
fn=cast(Any, fn),
|
||||
description=description or fn.__doc__,
|
||||
is_async=resolved_is_async,
|
||||
accepts_context=resolved_accepts_context,
|
||||
),
|
||||
)
|
||||
|
||||
if fn is not None:
|
||||
return decorator(fn)
|
||||
return decorator
|
||||
|
||||
|
||||
def build_registry(
|
||||
*specs: NodeSpec[Any, Any],
|
||||
) -> dict[str, SyncRegistryHandler]:
|
||||
return _build_registry(specs, export="sync")
|
||||
|
||||
|
||||
def build_async_registry(
|
||||
*specs: NodeSpec[Any, Any],
|
||||
) -> dict[str, AsyncRegistryHandler]:
|
||||
return _build_registry(specs, export="async")
|
||||
|
||||
|
||||
@overload
|
||||
def _build_registry(
|
||||
specs: tuple[NodeSpec[Any, Any], ...],
|
||||
*,
|
||||
export: Literal["sync"],
|
||||
) -> dict[str, SyncRegistryHandler]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def _build_registry(
|
||||
specs: tuple[NodeSpec[Any, Any], ...],
|
||||
*,
|
||||
export: Literal["async"],
|
||||
) -> dict[str, AsyncRegistryHandler]: ...
|
||||
|
||||
|
||||
def _build_registry(
|
||||
specs: tuple[NodeSpec[Any, Any], ...],
|
||||
*,
|
||||
export: Literal["sync", "async"],
|
||||
) -> dict[str, Any]:
|
||||
if export == "sync":
|
||||
return {spec.name: spec.to_registry_handler() for spec in specs}
|
||||
if export == "async":
|
||||
return {spec.name: spec.to_async_registry_handler() for spec in specs}
|
||||
raise ValueError(f"unknown registry export mode {export!r}")
|
||||
|
||||
@@ -1,343 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from .errors import WorkflowExecutionError
|
||||
from .foreach_ops import step_foreach
|
||||
from .flow_ops import advance_frame, append_step_result_trace, finalize_run
|
||||
from .frame_ops import collapse_completed_frames
|
||||
from .interrupt_ops import resume_interrupt
|
||||
from .model import (
|
||||
ConditionNode,
|
||||
ForeachNode,
|
||||
InterruptNode,
|
||||
JoinNode,
|
||||
NodeUse,
|
||||
Workflow,
|
||||
)
|
||||
from .node_exec import (
|
||||
AsyncNodeHandler,
|
||||
NodeHandler,
|
||||
coerce_node_result,
|
||||
execute_node_use,
|
||||
execute_node_use_async,
|
||||
)
|
||||
from .run_factory import create_run_state
|
||||
from .run_state import (
|
||||
FrameStatus,
|
||||
RunState,
|
||||
RunStatus,
|
||||
)
|
||||
from .schema_tools import validate_payload_against_schema
|
||||
from .step_handlers import (
|
||||
handle_condition_step,
|
||||
handle_interrupt_step,
|
||||
handle_join_step,
|
||||
)
|
||||
from .tokens import END
|
||||
from .workflow_index import WorkflowIndex, build_workflow_index
|
||||
|
||||
__all__ = [
|
||||
"AsyncNodeHandler",
|
||||
"NodeHandler",
|
||||
"coerce_node_result",
|
||||
"execute_workflow_async",
|
||||
"execute_workflow",
|
||||
"resume_workflow_async",
|
||||
"resume_workflow",
|
||||
"step_workflow_async",
|
||||
"step_workflow",
|
||||
]
|
||||
|
||||
|
||||
def _prepare_new_run(workflow: Workflow, workflow_input: dict[str, Any]) -> RunState:
|
||||
run = create_run_state(workflow, workflow_input)
|
||||
workflow.validate_structure().raise_for_errors()
|
||||
validate_payload_against_schema(
|
||||
workflow.input_schema, workflow_input, "workflow input"
|
||||
)
|
||||
return run
|
||||
|
||||
|
||||
def _prepare_resume(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
*,
|
||||
resume_payload: dict[str, Any] | None,
|
||||
resume_outcome: str,
|
||||
) -> WorkflowIndex | None:
|
||||
if run.workflow_name != workflow.name:
|
||||
raise WorkflowExecutionError(
|
||||
f"run state belongs to workflow {run.workflow_name!r}, not {workflow.name!r}"
|
||||
)
|
||||
|
||||
if run.current_frame_id is None:
|
||||
raise WorkflowExecutionError("run has no current frame")
|
||||
collapse_completed_frames(run)
|
||||
|
||||
if run.current_node_id is None:
|
||||
raise WorkflowExecutionError("run has no current node")
|
||||
|
||||
if run.status == RunStatus.COMPLETED:
|
||||
return None
|
||||
|
||||
index = build_workflow_index(workflow)
|
||||
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
if resume_payload is None:
|
||||
return None
|
||||
resume_interrupt(
|
||||
workflow,
|
||||
run,
|
||||
index=index,
|
||||
resume_payload=resume_payload,
|
||||
resume_outcome=resume_outcome,
|
||||
)
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id == END:
|
||||
return None
|
||||
|
||||
run.status = RunStatus.RUNNING
|
||||
run.error = None
|
||||
run.current_frame().status = FrameStatus.RUNNING
|
||||
return index
|
||||
|
||||
|
||||
def _prepare_step(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
index: WorkflowIndex | None,
|
||||
) -> tuple[WorkflowIndex, object] | None:
|
||||
if run.current_frame_id is None:
|
||||
raise WorkflowExecutionError("run has no current frame")
|
||||
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id is None or run.current_node_id == END:
|
||||
return None
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
return None
|
||||
|
||||
if run.status == RunStatus.PENDING:
|
||||
run.status = RunStatus.RUNNING
|
||||
run.error = None
|
||||
|
||||
resolved_index = index or build_workflow_index(workflow)
|
||||
frame = run.current_frame()
|
||||
if frame.status == FrameStatus.PENDING:
|
||||
frame.status = FrameStatus.RUNNING
|
||||
step = resolved_index.nodes_by_id[frame.node_id]
|
||||
return resolved_index, step
|
||||
|
||||
|
||||
def _complete_step(
|
||||
*,
|
||||
run: RunState,
|
||||
index: WorkflowIndex,
|
||||
outcome: str,
|
||||
frame_id: str,
|
||||
node_id: str,
|
||||
step_type: str,
|
||||
step_result: Any,
|
||||
) -> RunState:
|
||||
next_node_id = index.next_node_id(node_id, outcome)
|
||||
|
||||
append_step_result_trace(
|
||||
run,
|
||||
frame_id=frame_id,
|
||||
node_id=node_id,
|
||||
step_type=step_type,
|
||||
next_node_id=next_node_id,
|
||||
result=step_result,
|
||||
)
|
||||
advance_frame(
|
||||
run,
|
||||
run.frames[frame_id],
|
||||
outcome=outcome,
|
||||
next_node_id=next_node_id,
|
||||
)
|
||||
return run
|
||||
|
||||
|
||||
def execute_workflow(
|
||||
workflow: Workflow,
|
||||
workflow_input: dict[str, Any],
|
||||
registry: Mapping[str, NodeHandler],
|
||||
) -> RunState:
|
||||
run = create_run_state(workflow, workflow_input)
|
||||
|
||||
try:
|
||||
run = _prepare_new_run(workflow, workflow_input)
|
||||
return resume_workflow(workflow, run, registry)
|
||||
except Exception as exc:
|
||||
run.status = RunStatus.FAILED
|
||||
run.error = str(exc)
|
||||
raise
|
||||
|
||||
|
||||
async def execute_workflow_async(
|
||||
workflow: Workflow,
|
||||
workflow_input: dict[str, Any],
|
||||
registry: Mapping[str, AsyncNodeHandler],
|
||||
) -> RunState:
|
||||
run = create_run_state(workflow, workflow_input)
|
||||
|
||||
try:
|
||||
run = _prepare_new_run(workflow, workflow_input)
|
||||
return await resume_workflow_async(workflow, run, registry)
|
||||
except Exception as exc:
|
||||
run.status = RunStatus.FAILED
|
||||
run.error = str(exc)
|
||||
raise
|
||||
|
||||
|
||||
def resume_workflow(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, NodeHandler],
|
||||
*,
|
||||
resume_payload: dict[str, Any] | None = None,
|
||||
resume_outcome: str = "submitted",
|
||||
) -> RunState:
|
||||
index = _prepare_resume(
|
||||
workflow,
|
||||
run,
|
||||
resume_payload=resume_payload,
|
||||
resume_outcome=resume_outcome,
|
||||
)
|
||||
if index is None:
|
||||
if run.current_node_id == END:
|
||||
return finalize_run(workflow, run)
|
||||
return run
|
||||
|
||||
while True:
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id == END:
|
||||
break
|
||||
step_workflow(
|
||||
workflow,
|
||||
run,
|
||||
registry,
|
||||
index=index,
|
||||
)
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
return run
|
||||
|
||||
return finalize_run(workflow, run)
|
||||
|
||||
|
||||
async def resume_workflow_async(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, AsyncNodeHandler],
|
||||
*,
|
||||
resume_payload: dict[str, Any] | None = None,
|
||||
resume_outcome: str = "submitted",
|
||||
) -> RunState:
|
||||
index = _prepare_resume(
|
||||
workflow,
|
||||
run,
|
||||
resume_payload=resume_payload,
|
||||
resume_outcome=resume_outcome,
|
||||
)
|
||||
if index is None:
|
||||
if run.current_node_id == END:
|
||||
return finalize_run(workflow, run)
|
||||
return run
|
||||
|
||||
while True:
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id == END:
|
||||
break
|
||||
await step_workflow_async(
|
||||
workflow,
|
||||
run,
|
||||
registry,
|
||||
index=index,
|
||||
)
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
return run
|
||||
|
||||
return finalize_run(workflow, run)
|
||||
|
||||
|
||||
def step_workflow(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, NodeHandler],
|
||||
*,
|
||||
index: WorkflowIndex | None = None,
|
||||
) -> RunState:
|
||||
prepared = _prepare_step(workflow, run, index)
|
||||
if prepared is None:
|
||||
return run
|
||||
index, step = prepared
|
||||
frame = run.current_frame()
|
||||
|
||||
if isinstance(step, NodeUse):
|
||||
node_def = index.node_defs[step.node]
|
||||
step_result = execute_node_use(workflow, run, step, node_def, registry)
|
||||
elif isinstance(step, ConditionNode):
|
||||
step_result = handle_condition_step(run, step)
|
||||
elif isinstance(step, JoinNode):
|
||||
step_result = handle_join_step()
|
||||
elif isinstance(step, InterruptNode):
|
||||
return handle_interrupt_step(run, step)
|
||||
elif isinstance(step, ForeachNode):
|
||||
return step_foreach(workflow, run, step, index)
|
||||
else:
|
||||
raise WorkflowExecutionError(
|
||||
f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"
|
||||
)
|
||||
|
||||
return _complete_step(
|
||||
run=run,
|
||||
index=index,
|
||||
outcome=step_result.outcome,
|
||||
frame_id=frame.id,
|
||||
node_id=frame.node_id,
|
||||
step_type=step.type,
|
||||
step_result=step_result,
|
||||
)
|
||||
|
||||
|
||||
async def step_workflow_async(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, AsyncNodeHandler],
|
||||
*,
|
||||
index: WorkflowIndex | None = None,
|
||||
) -> RunState:
|
||||
prepared = _prepare_step(workflow, run, index)
|
||||
if prepared is None:
|
||||
return run
|
||||
index, step = prepared
|
||||
frame = run.current_frame()
|
||||
|
||||
if isinstance(step, NodeUse):
|
||||
node_def = index.node_defs[step.node]
|
||||
step_result = await execute_node_use_async(
|
||||
workflow, run, step, node_def, registry
|
||||
)
|
||||
elif isinstance(step, ConditionNode):
|
||||
step_result = handle_condition_step(run, step)
|
||||
elif isinstance(step, JoinNode):
|
||||
step_result = handle_join_step()
|
||||
elif isinstance(step, InterruptNode):
|
||||
return handle_interrupt_step(run, step)
|
||||
elif isinstance(step, ForeachNode):
|
||||
return step_foreach(workflow, run, step, index)
|
||||
else:
|
||||
raise WorkflowExecutionError(
|
||||
f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"
|
||||
)
|
||||
|
||||
return _complete_step(
|
||||
run=run,
|
||||
index=index,
|
||||
outcome=step_result.outcome,
|
||||
frame_id=frame.id,
|
||||
node_id=frame.node_id,
|
||||
step_type=step.type,
|
||||
step_result=step_result,
|
||||
)
|
||||
@@ -0,0 +1,32 @@
|
||||
from wf_core.errors import WorkflowExecutionError
|
||||
from wf_core.node_exec import (
|
||||
AsyncNodeHandler,
|
||||
NodeHandler,
|
||||
coerce_node_result,
|
||||
execute_node_use,
|
||||
execute_node_use_async,
|
||||
)
|
||||
|
||||
from .engine import (
|
||||
execute_workflow,
|
||||
execute_workflow_async,
|
||||
resume_workflow,
|
||||
resume_workflow_async,
|
||||
)
|
||||
from .step import complete_step, step_workflow, step_workflow_async
|
||||
|
||||
__all__ = [
|
||||
"AsyncNodeHandler",
|
||||
"NodeHandler",
|
||||
"WorkflowExecutionError",
|
||||
"coerce_node_result",
|
||||
"complete_step",
|
||||
"execute_node_use",
|
||||
"execute_node_use_async",
|
||||
"execute_workflow",
|
||||
"execute_workflow_async",
|
||||
"resume_workflow",
|
||||
"resume_workflow_async",
|
||||
"step_workflow",
|
||||
"step_workflow_async",
|
||||
]
|
||||
@@ -0,0 +1,121 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from wf_core.flow_ops import finalize_run
|
||||
from wf_core.frame_ops import collapse_completed_frames
|
||||
from wf_core.model import Workflow
|
||||
from wf_core.node_exec import AsyncNodeHandler, NodeHandler
|
||||
from wf_core.run_factory import create_run_state
|
||||
from wf_core.run_state import RunState, RunStatus
|
||||
from wf_core.tokens import END
|
||||
|
||||
from .preparation import prepare_new_run, prepare_resume
|
||||
from .step import step_workflow, step_workflow_async
|
||||
|
||||
|
||||
def execute_workflow(
|
||||
workflow: Workflow,
|
||||
workflow_input: dict[str, Any],
|
||||
registry: Mapping[str, NodeHandler],
|
||||
) -> RunState:
|
||||
"""Create a run and execute a workflow synchronously until it stops."""
|
||||
run = create_run_state(workflow, workflow_input)
|
||||
|
||||
try:
|
||||
run = prepare_new_run(workflow, workflow_input)
|
||||
return resume_workflow(workflow, run, registry)
|
||||
except Exception as exc:
|
||||
run.status = RunStatus.FAILED
|
||||
run.error = str(exc)
|
||||
raise
|
||||
|
||||
|
||||
async def execute_workflow_async(
|
||||
workflow: Workflow,
|
||||
workflow_input: dict[str, Any],
|
||||
registry: Mapping[str, AsyncNodeHandler],
|
||||
) -> RunState:
|
||||
"""Create a run and execute a workflow asynchronously until it stops."""
|
||||
run = create_run_state(workflow, workflow_input)
|
||||
|
||||
try:
|
||||
run = prepare_new_run(workflow, workflow_input)
|
||||
return await resume_workflow_async(workflow, run, registry)
|
||||
except Exception as exc:
|
||||
run.status = RunStatus.FAILED
|
||||
run.error = str(exc)
|
||||
raise
|
||||
|
||||
|
||||
def resume_workflow(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, NodeHandler],
|
||||
*,
|
||||
resume_payload: dict[str, Any] | None = None,
|
||||
resume_outcome: str = "submitted",
|
||||
) -> RunState:
|
||||
"""Resume a synchronous run from its current state."""
|
||||
index = prepare_resume(
|
||||
workflow,
|
||||
run,
|
||||
resume_payload=resume_payload,
|
||||
resume_outcome=resume_outcome,
|
||||
)
|
||||
if index is None:
|
||||
if run.current_node_id == END:
|
||||
return finalize_run(workflow, run)
|
||||
return run
|
||||
|
||||
while True:
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id == END:
|
||||
break
|
||||
step_workflow(
|
||||
workflow,
|
||||
run,
|
||||
registry,
|
||||
index=index,
|
||||
)
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
return run
|
||||
|
||||
return finalize_run(workflow, run)
|
||||
|
||||
|
||||
async def resume_workflow_async(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, AsyncNodeHandler],
|
||||
*,
|
||||
resume_payload: dict[str, Any] | None = None,
|
||||
resume_outcome: str = "submitted",
|
||||
) -> RunState:
|
||||
"""Resume an async run from its current state."""
|
||||
index = prepare_resume(
|
||||
workflow,
|
||||
run,
|
||||
resume_payload=resume_payload,
|
||||
resume_outcome=resume_outcome,
|
||||
)
|
||||
if index is None:
|
||||
if run.current_node_id == END:
|
||||
return finalize_run(workflow, run)
|
||||
return run
|
||||
|
||||
while True:
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id == END:
|
||||
break
|
||||
await step_workflow_async(
|
||||
workflow,
|
||||
run,
|
||||
registry,
|
||||
index=index,
|
||||
)
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
return run
|
||||
|
||||
return finalize_run(workflow, run)
|
||||
@@ -0,0 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from wf_core.errors import WorkflowExecutionError
|
||||
from wf_core.frame_ops import collapse_completed_frames
|
||||
from wf_core.interrupt_ops import resume_interrupt
|
||||
from wf_core.model import Workflow
|
||||
from wf_core.run_factory import create_run_state
|
||||
from wf_core.run_state import FrameStatus, RunState, RunStatus
|
||||
from wf_core.schema_tools import validate_payload_against_schema
|
||||
from wf_core.tokens import END
|
||||
from wf_core.workflow_index import WorkflowIndex, build_workflow_index
|
||||
|
||||
|
||||
def prepare_new_run(workflow: Workflow, workflow_input: dict[str, Any]) -> RunState:
|
||||
"""Create and validate a fresh run state for a workflow invocation."""
|
||||
run = create_run_state(workflow, workflow_input)
|
||||
workflow.validate_structure().raise_for_errors()
|
||||
validate_payload_against_schema(
|
||||
workflow.input_schema, workflow_input, "workflow input"
|
||||
)
|
||||
return run
|
||||
|
||||
|
||||
def prepare_resume(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
*,
|
||||
resume_payload: dict[str, Any] | None,
|
||||
resume_outcome: str,
|
||||
) -> WorkflowIndex | None:
|
||||
"""Validate and normalize a run state before resume execution."""
|
||||
if run.workflow_name != workflow.name:
|
||||
raise WorkflowExecutionError(
|
||||
f"run state belongs to workflow {run.workflow_name!r}, not {workflow.name!r}"
|
||||
)
|
||||
|
||||
if run.current_frame_id is None:
|
||||
raise WorkflowExecutionError("run has no current frame")
|
||||
collapse_completed_frames(run)
|
||||
|
||||
if run.current_node_id is None:
|
||||
raise WorkflowExecutionError("run has no current node")
|
||||
|
||||
if run.status == RunStatus.COMPLETED:
|
||||
return None
|
||||
|
||||
index = build_workflow_index(workflow)
|
||||
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
if resume_payload is None:
|
||||
return None
|
||||
resume_interrupt(
|
||||
workflow,
|
||||
run,
|
||||
index=index,
|
||||
resume_payload=resume_payload,
|
||||
resume_outcome=resume_outcome,
|
||||
)
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id == END:
|
||||
return None
|
||||
|
||||
run.status = RunStatus.RUNNING
|
||||
run.error = None
|
||||
run.current_frame().status = FrameStatus.RUNNING
|
||||
return index
|
||||
|
||||
|
||||
def prepare_step(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
index: WorkflowIndex | None,
|
||||
) -> tuple[WorkflowIndex, object] | None:
|
||||
"""Resolve the next executable workflow step for the current run frame."""
|
||||
if run.current_frame_id is None:
|
||||
raise WorkflowExecutionError("run has no current frame")
|
||||
|
||||
collapse_completed_frames(run)
|
||||
if run.current_node_id is None or run.current_node_id == END:
|
||||
return None
|
||||
if run.status == RunStatus.INTERRUPTED:
|
||||
return None
|
||||
|
||||
if run.status == RunStatus.PENDING:
|
||||
run.status = RunStatus.RUNNING
|
||||
run.error = None
|
||||
|
||||
resolved_index = index or build_workflow_index(workflow)
|
||||
frame = run.current_frame()
|
||||
if frame.status == FrameStatus.PENDING:
|
||||
frame.status = FrameStatus.RUNNING
|
||||
step = resolved_index.nodes_by_id[frame.node_id]
|
||||
return resolved_index, step
|
||||
@@ -0,0 +1,145 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from wf_core.errors import WorkflowExecutionError
|
||||
from wf_core.flow_ops import advance_frame, append_step_result_trace
|
||||
from wf_core.foreach_ops import step_foreach
|
||||
from wf_core.model import (
|
||||
ConditionNode,
|
||||
ForeachNode,
|
||||
InterruptNode,
|
||||
JoinNode,
|
||||
NodeUse,
|
||||
Workflow,
|
||||
)
|
||||
from wf_core.node_exec import (
|
||||
AsyncNodeHandler,
|
||||
NodeHandler,
|
||||
execute_node_use,
|
||||
execute_node_use_async,
|
||||
)
|
||||
from wf_core.run_state import RunState
|
||||
from wf_core.step_handlers import (
|
||||
handle_condition_step,
|
||||
handle_interrupt_step,
|
||||
handle_join_step,
|
||||
)
|
||||
from wf_core.workflow_index import WorkflowIndex
|
||||
|
||||
from .preparation import prepare_step
|
||||
|
||||
|
||||
def complete_step(
|
||||
*,
|
||||
run: RunState,
|
||||
index: WorkflowIndex,
|
||||
outcome: str,
|
||||
frame_id: str,
|
||||
node_id: str,
|
||||
step_type: str,
|
||||
step_result: Any,
|
||||
) -> RunState:
|
||||
"""Record a completed step and advance the active frame."""
|
||||
next_node_id = index.next_node_id(node_id, outcome)
|
||||
|
||||
append_step_result_trace(
|
||||
run,
|
||||
frame_id=frame_id,
|
||||
node_id=node_id,
|
||||
step_type=step_type,
|
||||
next_node_id=next_node_id,
|
||||
result=step_result,
|
||||
)
|
||||
advance_frame(
|
||||
run,
|
||||
run.frames[frame_id],
|
||||
outcome=outcome,
|
||||
next_node_id=next_node_id,
|
||||
)
|
||||
return run
|
||||
|
||||
|
||||
def step_workflow(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, NodeHandler],
|
||||
*,
|
||||
index: WorkflowIndex | None = None,
|
||||
) -> RunState:
|
||||
"""Execute at most one synchronous workflow step."""
|
||||
prepared = prepare_step(workflow, run, index)
|
||||
if prepared is None:
|
||||
return run
|
||||
index, step = prepared
|
||||
frame = run.current_frame()
|
||||
|
||||
if isinstance(step, NodeUse):
|
||||
node_def = index.node_defs[step.node]
|
||||
step_result = execute_node_use(workflow, run, step, node_def, registry)
|
||||
elif isinstance(step, ConditionNode):
|
||||
step_result = handle_condition_step(run, step)
|
||||
elif isinstance(step, JoinNode):
|
||||
step_result = handle_join_step()
|
||||
elif isinstance(step, InterruptNode):
|
||||
return handle_interrupt_step(run, step)
|
||||
elif isinstance(step, ForeachNode):
|
||||
return step_foreach(workflow, run, step, index)
|
||||
else:
|
||||
raise WorkflowExecutionError(
|
||||
f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"
|
||||
)
|
||||
|
||||
return complete_step(
|
||||
run=run,
|
||||
index=index,
|
||||
outcome=step_result.outcome,
|
||||
frame_id=frame.id,
|
||||
node_id=frame.node_id,
|
||||
step_type=step.type,
|
||||
step_result=step_result,
|
||||
)
|
||||
|
||||
|
||||
async def step_workflow_async(
|
||||
workflow: Workflow,
|
||||
run: RunState,
|
||||
registry: Mapping[str, AsyncNodeHandler],
|
||||
*,
|
||||
index: WorkflowIndex | None = None,
|
||||
) -> RunState:
|
||||
"""Execute at most one async workflow step."""
|
||||
prepared = prepare_step(workflow, run, index)
|
||||
if prepared is None:
|
||||
return run
|
||||
index, step = prepared
|
||||
frame = run.current_frame()
|
||||
|
||||
if isinstance(step, NodeUse):
|
||||
node_def = index.node_defs[step.node]
|
||||
step_result = await execute_node_use_async(
|
||||
workflow, run, step, node_def, registry
|
||||
)
|
||||
elif isinstance(step, ConditionNode):
|
||||
step_result = handle_condition_step(run, step)
|
||||
elif isinstance(step, JoinNode):
|
||||
step_result = handle_join_step()
|
||||
elif isinstance(step, InterruptNode):
|
||||
return handle_interrupt_step(run, step)
|
||||
elif isinstance(step, ForeachNode):
|
||||
return step_foreach(workflow, run, step, index)
|
||||
else:
|
||||
raise WorkflowExecutionError(
|
||||
f"unsupported step type {getattr(step, 'type', type(step).__name__)!r}"
|
||||
)
|
||||
|
||||
return complete_step(
|
||||
run=run,
|
||||
index=index,
|
||||
outcome=step_result.outcome,
|
||||
frame_id=frame.id,
|
||||
node_id=frame.node_id,
|
||||
step_type=step.type,
|
||||
step_result=step_result,
|
||||
)
|
||||
Reference in New Issue
Block a user