and thats a server we can use
This commit is contained in:
@@ -0,0 +1,42 @@
|
||||
from .builder import WorkflowBuilder
|
||||
from .catalog import NodeCatalog, NodeCatalogEntry
|
||||
from .conditions import context, exists, expr, input, state
|
||||
from .mapping import bind_fields, bind_state, merge_maps
|
||||
from .paths import GraphPath, context_path, graph_path, input_path, state_path
|
||||
from .spec import (
|
||||
AsyncRegistryHandler,
|
||||
NodeReturn,
|
||||
NodeSpec,
|
||||
SyncRegistryHandler,
|
||||
build_async_registry,
|
||||
build_registry,
|
||||
node,
|
||||
)
|
||||
from .subgraph import subgraph_node
|
||||
|
||||
__all__ = [
|
||||
"NodeCatalog",
|
||||
"NodeCatalogEntry",
|
||||
"GraphPath",
|
||||
"NodeReturn",
|
||||
"NodeSpec",
|
||||
"AsyncRegistryHandler",
|
||||
"SyncRegistryHandler",
|
||||
"WorkflowBuilder",
|
||||
"bind_fields",
|
||||
"build_async_registry",
|
||||
"build_registry",
|
||||
"bind_state",
|
||||
"merge_maps",
|
||||
"context",
|
||||
"context_path",
|
||||
"expr",
|
||||
"exists",
|
||||
"graph_path",
|
||||
"input",
|
||||
"input_path",
|
||||
"node",
|
||||
"state",
|
||||
"state_path",
|
||||
"subgraph_node",
|
||||
]
|
||||
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal, TypeAlias
|
||||
|
||||
from wf_core import (
|
||||
ConditionNode,
|
||||
Edge,
|
||||
ForeachNode,
|
||||
InterruptNode,
|
||||
NodeUse,
|
||||
SchemaRef,
|
||||
StateSchema,
|
||||
Workflow,
|
||||
)
|
||||
from wf_core.model import Condition as CoreCondition
|
||||
|
||||
from .conditions import Expr, compile_condition
|
||||
from .mapping import PathArg
|
||||
from .paths import GraphPath
|
||||
from .spec import NodeSpec
|
||||
|
||||
StepRef: TypeAlias = str | NodeUse | ConditionNode | ForeachNode | InterruptNode
|
||||
MapArg: TypeAlias = Mapping[Any, Any]
|
||||
|
||||
|
||||
def _coerce_path(value: object) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, GraphPath):
|
||||
return value.value
|
||||
raise TypeError(f"unsupported graph path value {value!r}")
|
||||
|
||||
|
||||
def _normalize_mapping(
|
||||
mapping: MapArg | None,
|
||||
) -> dict[str, str]:
|
||||
if mapping is None:
|
||||
return {}
|
||||
return {
|
||||
_coerce_path(source): _coerce_path(destination)
|
||||
for source, destination in mapping.items()
|
||||
}
|
||||
|
||||
|
||||
def _step_id(ref: StepRef) -> str:
|
||||
if isinstance(ref, str):
|
||||
return ref
|
||||
return ref.id
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class WorkflowBuilder:
|
||||
name: str
|
||||
input_schema: SchemaRef
|
||||
state_schema: StateSchema
|
||||
output_schema: SchemaRef
|
||||
start: str
|
||||
node_specs: dict[str, NodeSpec[Any, Any]] = field(default_factory=dict)
|
||||
nodes: list[Any] = field(default_factory=list)
|
||||
edges: list[Edge] = field(default_factory=list)
|
||||
|
||||
def use(
|
||||
self,
|
||||
spec: NodeSpec[Any, Any],
|
||||
*,
|
||||
id: str,
|
||||
in_map: MapArg | None = None,
|
||||
out_map: MapArg | None = None,
|
||||
desc: str | None = None,
|
||||
) -> NodeUse:
|
||||
self.node_specs[spec.name] = spec
|
||||
node = NodeUse(
|
||||
id=id,
|
||||
type="node",
|
||||
node=spec.name,
|
||||
desc=desc or spec.description,
|
||||
in_map=_normalize_mapping(in_map),
|
||||
out_map=_normalize_mapping(out_map),
|
||||
)
|
||||
self.nodes.append(node)
|
||||
return node
|
||||
|
||||
def condition(self, *, id: str, check: CoreCondition | Expr) -> ConditionNode:
|
||||
node = ConditionNode(
|
||||
id=id,
|
||||
type="condition",
|
||||
check=compile_condition(check),
|
||||
)
|
||||
self.nodes.append(node)
|
||||
return node
|
||||
|
||||
def foreach(
|
||||
self,
|
||||
*,
|
||||
id: str,
|
||||
over: PathArg,
|
||||
as_: str,
|
||||
mode: Literal["serial", "parallel"] = "serial",
|
||||
on_item_error: Literal["fail", "collect", "skip"] = "fail",
|
||||
) -> ForeachNode:
|
||||
node = ForeachNode.model_validate(
|
||||
{
|
||||
"id": id,
|
||||
"type": "foreach",
|
||||
"over": _coerce_path(over),
|
||||
"as": as_,
|
||||
"mode": mode,
|
||||
"on_item_error": on_item_error,
|
||||
}
|
||||
)
|
||||
self.nodes.append(node)
|
||||
return node
|
||||
|
||||
def interrupt(
|
||||
self,
|
||||
*,
|
||||
id: str,
|
||||
kind: str,
|
||||
request_map: MapArg | None = None,
|
||||
out_map: MapArg | None = None,
|
||||
outcomes: list[str] | None = None,
|
||||
) -> InterruptNode:
|
||||
node = InterruptNode(
|
||||
id=id,
|
||||
type="interrupt",
|
||||
kind=kind,
|
||||
request_map=_normalize_mapping(request_map),
|
||||
out_map=_normalize_mapping(out_map),
|
||||
outcomes=outcomes or ["submitted"],
|
||||
)
|
||||
self.nodes.append(node)
|
||||
return node
|
||||
|
||||
def connect(self, from_: StepRef, outcome: str, to: StepRef) -> None:
|
||||
self.edges.append(
|
||||
Edge.model_validate(
|
||||
{"from": _step_id(from_), "outcome": outcome, "to": _step_id(to)}
|
||||
)
|
||||
)
|
||||
|
||||
def compile(self) -> Workflow:
|
||||
node_defs = [spec.to_node_def() for spec in self.node_specs.values()]
|
||||
return Workflow(
|
||||
name=self.name,
|
||||
input_schema=self.input_schema,
|
||||
state_schema=self.state_schema,
|
||||
output_schema=self.output_schema,
|
||||
node_defs=node_defs,
|
||||
start=self.start,
|
||||
nodes=self.nodes,
|
||||
edges=self.edges,
|
||||
)
|
||||
@@ -0,0 +1,51 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from .spec import NodeSpec
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class NodeCatalogEntry:
|
||||
name: str
|
||||
description: str | None
|
||||
outcomes: tuple[str, ...]
|
||||
input_schema: dict[str, Any]
|
||||
output_schema: dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_spec(cls, spec: NodeSpec[Any, Any]) -> "NodeCatalogEntry":
|
||||
return cls(
|
||||
name=spec.name,
|
||||
description=spec.description,
|
||||
outcomes=spec.outcomes,
|
||||
input_schema=spec.input_model.model_json_schema(),
|
||||
output_schema=spec.output_model.model_json_schema(),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class NodeCatalog:
|
||||
specs: dict[str, NodeSpec[Any, Any]]
|
||||
|
||||
@classmethod
|
||||
def from_specs(cls, *specs: NodeSpec[Any, Any]) -> "NodeCatalog":
|
||||
return cls(specs={spec.name: spec for spec in specs})
|
||||
|
||||
def entries(self) -> list[NodeCatalogEntry]:
|
||||
return [NodeCatalogEntry.from_spec(spec) for spec in self.specs.values()]
|
||||
|
||||
def as_mcp_payload(self) -> dict[str, Any]:
|
||||
return {
|
||||
"nodes": [
|
||||
{
|
||||
"name": entry.name,
|
||||
"description": entry.description,
|
||||
"outcomes": list(entry.outcomes),
|
||||
"input_schema": entry.input_schema,
|
||||
"output_schema": entry.output_schema,
|
||||
}
|
||||
for entry in self.entries()
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
from wf_core.model import (
|
||||
BinaryCondition,
|
||||
Condition,
|
||||
ExistsCondition,
|
||||
LiteralOperand,
|
||||
NotCondition,
|
||||
PathOperand,
|
||||
VariadicCondition,
|
||||
)
|
||||
from .paths import GraphPath, context_path, input_path, state_path
|
||||
|
||||
|
||||
def _operand(value: object) -> PathOperand | LiteralOperand:
|
||||
if isinstance(value, PathExpr):
|
||||
return PathOperand(path=value.path)
|
||||
if isinstance(value, GraphPath):
|
||||
return PathOperand(path=value.value)
|
||||
return LiteralOperand(value=value)
|
||||
|
||||
|
||||
def _path_str(value: PathExpr | GraphPath) -> str:
|
||||
if isinstance(value, PathExpr):
|
||||
return value.path
|
||||
return value.value
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Expr:
|
||||
condition: Condition
|
||||
|
||||
def __and__(self, other: object) -> Expr:
|
||||
if not isinstance(other, Expr):
|
||||
return NotImplemented
|
||||
return Expr(VariadicCondition(op="and", args=[self.condition, other.condition]))
|
||||
|
||||
def __or__(self, other: object) -> Expr:
|
||||
if not isinstance(other, Expr):
|
||||
return NotImplemented
|
||||
return Expr(VariadicCondition(op="or", args=[self.condition, other.condition]))
|
||||
|
||||
def __invert__(self) -> Expr:
|
||||
return Expr(NotCondition(op="not", arg=self.condition))
|
||||
|
||||
def to_condition(self) -> Condition:
|
||||
return self.condition
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PathExpr:
|
||||
path: str
|
||||
|
||||
def _binary(self, op: Literal["eq", "ne", "gt", "lt"], other: object) -> Expr:
|
||||
return Expr(
|
||||
BinaryCondition(
|
||||
op=op,
|
||||
left=PathOperand(path=self.path),
|
||||
right=_operand(other),
|
||||
)
|
||||
)
|
||||
|
||||
def eq(self, other: object) -> Expr:
|
||||
return self._binary("eq", other)
|
||||
|
||||
def ne(self, other: object) -> Expr:
|
||||
return self._binary("ne", other)
|
||||
|
||||
def gt(self, other: object) -> Expr:
|
||||
return self._binary("gt", other)
|
||||
|
||||
def lt(self, other: object) -> Expr:
|
||||
return self._binary("lt", other)
|
||||
|
||||
def __eq__(self, other: object) -> Expr: # type: ignore[override] # ty: ignore[invalid-method-override]
|
||||
return self._binary("eq", other)
|
||||
|
||||
def __ne__(self, other: object) -> Expr: # type: ignore[override] # ty: ignore[invalid-method-override]
|
||||
return self._binary("ne", other)
|
||||
|
||||
def __gt__(self, other: object) -> Expr:
|
||||
return self.gt(other)
|
||||
|
||||
def __lt__(self, other: object) -> Expr:
|
||||
return self.lt(other)
|
||||
|
||||
|
||||
def expr(value: PathExpr | GraphPath) -> PathExpr:
|
||||
if isinstance(value, PathExpr):
|
||||
return value
|
||||
return PathExpr(path=value.value)
|
||||
|
||||
|
||||
def state(field: str) -> PathExpr:
|
||||
return expr(state_path(field))
|
||||
|
||||
|
||||
def input(field: str) -> PathExpr:
|
||||
return expr(input_path(field))
|
||||
|
||||
|
||||
def context(field: str) -> PathExpr:
|
||||
return expr(context_path(field))
|
||||
|
||||
|
||||
def exists(value: PathExpr | GraphPath) -> Expr:
|
||||
return Expr(ExistsCondition(op="exists", path=_path_str(value)))
|
||||
|
||||
|
||||
def compile_condition(value: Condition | Expr) -> Condition:
|
||||
if isinstance(value, Expr):
|
||||
return value.to_condition()
|
||||
return value
|
||||
@@ -0,0 +1,33 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TypeAlias
|
||||
|
||||
from .paths import GraphPath
|
||||
|
||||
PathArg: TypeAlias = str | GraphPath
|
||||
|
||||
|
||||
def normalize_path(path: PathArg) -> str:
|
||||
if isinstance(path, GraphPath):
|
||||
return path.value
|
||||
return path
|
||||
|
||||
|
||||
def bind_fields(**mapping: PathArg) -> dict[str, str]:
|
||||
return {
|
||||
normalize_path(source): destination for destination, source in mapping.items()
|
||||
}
|
||||
|
||||
|
||||
def bind_state(**mapping: PathArg) -> dict[str, str]:
|
||||
return {
|
||||
destination: normalize_path(target) for destination, target in mapping.items()
|
||||
}
|
||||
|
||||
|
||||
def merge_maps(*maps: Mapping[str, str]) -> dict[str, str]:
|
||||
merged: dict[str, str] = {}
|
||||
for mapping in maps:
|
||||
merged.update(mapping)
|
||||
return merged
|
||||
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GraphPath:
|
||||
value: str
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.value
|
||||
|
||||
|
||||
def graph_path(value: str) -> GraphPath:
|
||||
return GraphPath(value)
|
||||
|
||||
|
||||
def input_path(field: str) -> GraphPath:
|
||||
return GraphPath(f"input.{field}")
|
||||
|
||||
|
||||
def state_path(field: str) -> GraphPath:
|
||||
return GraphPath(f"state.{field}")
|
||||
|
||||
|
||||
def context_path(field: str) -> GraphPath:
|
||||
return GraphPath(f"context.{field}")
|
||||
@@ -0,0 +1,303 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from inspect import Parameter, iscoroutinefunction, signature
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
Any,
|
||||
Generic,
|
||||
Literal,
|
||||
TypeVar,
|
||||
cast,
|
||||
get_args,
|
||||
get_origin,
|
||||
get_type_hints,
|
||||
overload,
|
||||
)
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_core import NodeDef, RuntimeContext, SchemaRef
|
||||
|
||||
InputT = TypeVar("InputT", bound=BaseModel)
|
||||
OutputT = TypeVar("OutputT", bound=BaseModel)
|
||||
NodeCallable = Callable[[InputT, RuntimeContext], "NodeReturn[OutputT] | OutputT"]
|
||||
AsyncNodeCallable = Callable[
|
||||
[InputT, RuntimeContext], Awaitable["NodeReturn[OutputT] | OutputT"]
|
||||
]
|
||||
SyncRegistryHandler = Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]
|
||||
AsyncRegistryHandler = Callable[
|
||||
[dict[str, Any], RuntimeContext], Awaitable[dict[str, Any]]
|
||||
]
|
||||
|
||||
|
||||
def _schema_ref_for(model_type: type[BaseModel]) -> SchemaRef:
|
||||
return SchemaRef.model_validate(model_type.model_json_schema())
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class NodeReturn(Generic[OutputT]):
|
||||
outcome: str
|
||||
output: OutputT
|
||||
|
||||
|
||||
def _default_outcome(spec: "NodeSpec[Any, Any]") -> str:
|
||||
return spec.outcomes[0]
|
||||
|
||||
|
||||
def _coerce_registry_result(
|
||||
*,
|
||||
node_name: str,
|
||||
output_model: type[BaseModel],
|
||||
default_outcome: str,
|
||||
raw: NodeReturn[BaseModel] | BaseModel,
|
||||
) -> dict[str, Any]:
|
||||
if isinstance(raw, NodeReturn):
|
||||
if not isinstance(raw.output, output_model):
|
||||
raise TypeError(
|
||||
f"node {node_name!r} returned NodeReturn with unsupported output "
|
||||
f"{type(raw.output)!r}"
|
||||
)
|
||||
return {
|
||||
"outcome": raw.outcome,
|
||||
"output": raw.output.model_dump(),
|
||||
}
|
||||
if isinstance(raw, output_model):
|
||||
return {"outcome": default_outcome, "output": raw.model_dump()}
|
||||
raise TypeError(f"node {node_name!r} returned unsupported value {type(raw)!r}")
|
||||
|
||||
|
||||
def _is_basemodel_subclass(value: object) -> bool:
|
||||
return isinstance(value, type) and issubclass(value, BaseModel)
|
||||
|
||||
|
||||
def _infer_models(
|
||||
fn: Callable[..., object],
|
||||
) -> tuple[type[BaseModel], type[BaseModel]]:
|
||||
hints = get_type_hints(fn, include_extras=True)
|
||||
params = list(signature(fn).parameters.values())
|
||||
if len(params) < 2:
|
||||
raise TypeError("node function must accept at least (payload, ctx) parameters")
|
||||
|
||||
payload_param = params[0]
|
||||
ctx_param = params[1]
|
||||
|
||||
if payload_param.kind not in (
|
||||
Parameter.POSITIONAL_ONLY,
|
||||
Parameter.POSITIONAL_OR_KEYWORD,
|
||||
):
|
||||
raise TypeError("node payload parameter must be positional")
|
||||
if ctx_param.kind not in (
|
||||
Parameter.POSITIONAL_ONLY,
|
||||
Parameter.POSITIONAL_OR_KEYWORD,
|
||||
):
|
||||
raise TypeError("node context parameter must be positional")
|
||||
|
||||
input_model = hints.get(payload_param.name)
|
||||
if not _is_basemodel_subclass(input_model):
|
||||
raise TypeError("node payload annotation must be a pydantic BaseModel subclass")
|
||||
|
||||
ctx_type = hints.get(ctx_param.name)
|
||||
if ctx_type is not RuntimeContext:
|
||||
raise TypeError("node context annotation must be wf_core.RuntimeContext")
|
||||
|
||||
return_type = hints.get("return")
|
||||
if return_type is None:
|
||||
raise TypeError("node function must declare a return annotation")
|
||||
|
||||
if _is_basemodel_subclass(return_type):
|
||||
return cast(type[BaseModel], input_model), cast(type[BaseModel], return_type)
|
||||
|
||||
origin = get_origin(return_type)
|
||||
if origin is NodeReturn:
|
||||
args = get_args(return_type)
|
||||
if len(args) != 1 or not _is_basemodel_subclass(args[0]):
|
||||
raise TypeError(
|
||||
"NodeReturn return annotation must wrap a BaseModel subclass"
|
||||
)
|
||||
return cast(type[BaseModel], input_model), cast(type[BaseModel], args[0])
|
||||
|
||||
raise TypeError(
|
||||
"node return annotation must be a BaseModel subclass or NodeReturn[BaseModel]"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class NodeSpec(Generic[InputT, OutputT]):
|
||||
name: str
|
||||
input_model: type[InputT]
|
||||
output_model: type[OutputT]
|
||||
outcomes: tuple[str, ...]
|
||||
fn: NodeCallable[InputT, OutputT] | AsyncNodeCallable[InputT, OutputT]
|
||||
description: str | None = None
|
||||
is_async: bool = False
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
payload: InputT,
|
||||
ctx: RuntimeContext,
|
||||
) -> NodeReturn[OutputT] | OutputT | Awaitable[NodeReturn[OutputT] | OutputT]:
|
||||
return self.fn(payload, ctx)
|
||||
|
||||
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),
|
||||
outcomes=list(self.outcomes),
|
||||
)
|
||||
|
||||
def to_registry_handler(self) -> SyncRegistryHandler:
|
||||
if self.is_async:
|
||||
raise TypeError(
|
||||
f"node {self.name!r} is async and cannot be exported to the sync registry"
|
||||
)
|
||||
|
||||
def handler(payload: dict[str, Any], ctx: RuntimeContext) -> dict[str, Any]:
|
||||
parsed = self.input_model.model_validate(payload)
|
||||
raw = self.fn(parsed, ctx)
|
||||
return _coerce_registry_result(
|
||||
node_name=self.name,
|
||||
output_model=self.output_model,
|
||||
default_outcome=_default_outcome(self),
|
||||
raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
|
||||
)
|
||||
|
||||
return handler
|
||||
|
||||
def to_async_registry_handler(self) -> AsyncRegistryHandler:
|
||||
async def handler(
|
||||
payload: dict[str, Any],
|
||||
ctx: RuntimeContext,
|
||||
) -> dict[str, Any]:
|
||||
parsed = self.input_model.model_validate(payload)
|
||||
raw_result = self.fn(parsed, ctx)
|
||||
if self.is_async:
|
||||
raw = await cast(
|
||||
Awaitable[NodeReturn[OutputT] | OutputT],
|
||||
raw_result,
|
||||
)
|
||||
else:
|
||||
raw = cast(NodeReturn[OutputT] | OutputT, raw_result)
|
||||
return _coerce_registry_result(
|
||||
node_name=self.name,
|
||||
output_model=self.output_model,
|
||||
default_outcome=_default_outcome(self),
|
||||
raw=cast(NodeReturn[BaseModel] | BaseModel, raw),
|
||||
)
|
||||
|
||||
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
|
||||
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,
|
||||
),
|
||||
)
|
||||
|
||||
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}")
|
||||
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_core import RuntimeContext, Workflow, execute_workflow
|
||||
|
||||
from .spec import NodeSpec
|
||||
|
||||
InputT = TypeVar("InputT", bound=BaseModel)
|
||||
OutputT = TypeVar("OutputT", bound=BaseModel)
|
||||
|
||||
|
||||
def subgraph_node(
|
||||
*,
|
||||
name: str,
|
||||
workflow: Workflow,
|
||||
registry: Mapping[str, Any],
|
||||
input_model: type[InputT],
|
||||
output_model: type[OutputT],
|
||||
description: str | None = None,
|
||||
) -> NodeSpec[InputT, OutputT]:
|
||||
def run_subgraph(payload: InputT, ctx: RuntimeContext) -> OutputT:
|
||||
child_run = execute_workflow(
|
||||
workflow,
|
||||
payload.model_dump(),
|
||||
registry,
|
||||
)
|
||||
return output_model.model_validate(child_run.output)
|
||||
|
||||
return NodeSpec(
|
||||
name=name,
|
||||
input_model=input_model,
|
||||
output_model=output_model,
|
||||
outcomes=("ok",),
|
||||
fn=run_subgraph,
|
||||
description=description or f"Subgraph wrapper for {workflow.name}",
|
||||
is_async=False,
|
||||
)
|
||||
Reference in New Issue
Block a user