and more good stuff, fmt
This commit is contained in:
@@ -3,7 +3,8 @@ from __future__ import annotations
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
import re
|
||||
from typing import Any, Literal, TypeAlias, cast
|
||||
from typing import Any, Literal, TypeAlias, TypeGuard, cast
|
||||
import warnings
|
||||
|
||||
from wf_core import (
|
||||
ConditionNode,
|
||||
@@ -14,6 +15,8 @@ from wf_core import (
|
||||
SchemaRef,
|
||||
StateSchema,
|
||||
Workflow,
|
||||
RunState,
|
||||
execute_workflow,
|
||||
)
|
||||
from wf_core.errors import WorkflowExecutionError
|
||||
from wf_core.model import Condition as CoreCondition
|
||||
@@ -25,6 +28,7 @@ from .schemas import SchemaLike, StateSchemaLike, schema_ref_from, state_schema_
|
||||
from .spec import NodeSpec
|
||||
|
||||
StepRef: TypeAlias = str | NodeUse | ConditionNode | ForeachNode | InterruptNode
|
||||
BranchRef: TypeAlias = StepRef | NodeSpec[Any, Any]
|
||||
MapArg: TypeAlias = Mapping[Any, Any]
|
||||
|
||||
|
||||
@@ -53,6 +57,10 @@ def _step_id(ref: StepRef) -> str:
|
||||
return ref.id
|
||||
|
||||
|
||||
def _is_node_spec(ref: object) -> TypeGuard[NodeSpec[Any, Any]]:
|
||||
return isinstance(ref, NodeSpec)
|
||||
|
||||
|
||||
def _slug_id(value: str) -> str:
|
||||
slug = re.sub(r"[^0-9A-Za-z_]+", "_", value).strip("_").lower()
|
||||
return slug or "step"
|
||||
@@ -65,7 +73,9 @@ def _auto_input_map(
|
||||
state_schema: StateSchema,
|
||||
) -> dict[str, str]:
|
||||
return {
|
||||
_auto_source_path(field, input_schema=input_schema, state_schema=state_schema): field
|
||||
_auto_source_path(
|
||||
field, input_schema=input_schema, state_schema=state_schema
|
||||
): field
|
||||
for field in spec.input_model.model_json_schema().get("properties", {})
|
||||
}
|
||||
|
||||
@@ -165,6 +175,15 @@ class WorkflowBuilder:
|
||||
"""Export handlers for all node specs used by this builder."""
|
||||
return build_registry(*self.node_specs.values())
|
||||
|
||||
def execute(self, workflow_input: dict[str, Any]) -> RunState:
|
||||
"""Compile and execute this workflow with its used node registry.
|
||||
|
||||
This is intended for tests, examples, and local authoring loops. Production
|
||||
callers that need custom registries, persistence, or resume behavior should
|
||||
call wf_core execution functions directly.
|
||||
"""
|
||||
return execute_workflow(self.compile(), workflow_input, self.registry())
|
||||
|
||||
def condition(self, *, id: str, check: CoreCondition | Expr) -> ConditionNode:
|
||||
node = ConditionNode(
|
||||
id=id,
|
||||
@@ -223,6 +242,33 @@ class WorkflowBuilder:
|
||||
)
|
||||
)
|
||||
|
||||
def branch(
|
||||
self,
|
||||
from_: BranchRef,
|
||||
branches: Mapping[str, BranchRef],
|
||||
) -> dict[str, StepRef]:
|
||||
"""Connect multiple outcomes from one branch source.
|
||||
|
||||
Passing a NodeSpec creates a node use with auto-mapping and an auto id.
|
||||
Passing an existing step or id only wires edges. Empty branch maps are
|
||||
allowed but warn because they usually indicate an unfinished router.
|
||||
"""
|
||||
if not branches:
|
||||
warnings.warn(
|
||||
"WorkflowBuilder.branch called with no branches",
|
||||
UserWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
return {}
|
||||
|
||||
source = self.use(from_) if _is_node_spec(from_) else from_
|
||||
resolved_targets: dict[str, StepRef] = {}
|
||||
for outcome, target in branches.items():
|
||||
resolved = self.use(target) if _is_node_spec(target) else target
|
||||
self.connect(cast(StepRef, source), outcome, cast(StepRef, resolved))
|
||||
resolved_targets[outcome] = cast(StepRef, resolved)
|
||||
return resolved_targets
|
||||
|
||||
def compile(self) -> Workflow:
|
||||
if self.start is None:
|
||||
raise WorkflowExecutionError(
|
||||
|
||||
@@ -33,9 +33,7 @@ class PlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
|
||||
) -> NodeReturn[OutputT_co] | OutputT_co: ...
|
||||
|
||||
|
||||
NodeCallable = ContextNodeCallable[InputT, OutputT] | PlainNodeCallable[
|
||||
InputT, OutputT
|
||||
]
|
||||
NodeCallable = ContextNodeCallable[InputT, OutputT] | PlainNodeCallable[InputT, OutputT]
|
||||
|
||||
|
||||
class AsyncContextNodeCallable(Protocol[InputT_contra, OutputT_co]):
|
||||
@@ -55,9 +53,9 @@ class AsyncPlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
|
||||
) -> Awaitable[NodeReturn[OutputT_co] | OutputT_co]: ...
|
||||
|
||||
|
||||
AsyncNodeCallable = AsyncContextNodeCallable[
|
||||
InputT, OutputT
|
||||
] | AsyncPlainNodeCallable[InputT, OutputT]
|
||||
AsyncNodeCallable = (
|
||||
AsyncContextNodeCallable[InputT, OutputT] | AsyncPlainNodeCallable[InputT, OutputT]
|
||||
)
|
||||
|
||||
SyncRegistryHandler = Callable[[dict[str, Any], RuntimeContext], dict[str, Any]]
|
||||
AsyncRegistryHandler = Callable[
|
||||
|
||||
+2
-1
@@ -43,7 +43,8 @@ def parse_namespaced_tool_name(
|
||||
def is_admin_tool_name(proxy_name: str) -> bool:
|
||||
return proxy_name.startswith(f"{ADMIN_NAMESPACE}_")
|
||||
|
||||
|
||||
class LdaNamespace(Namespace):
|
||||
def __init__(self, prefix: str) -> None:
|
||||
super().__init__(prefix)
|
||||
self._name_prefix = f"{prefix}." # some good stuff
|
||||
self._name_prefix = f"{prefix}." # some good stuff
|
||||
|
||||
@@ -60,7 +60,9 @@ class TransparentProxyRuntime:
|
||||
if prompts_as_tools:
|
||||
self.server.add_transform(PromptsAsTools(self.server))
|
||||
if search_tools:
|
||||
self.server.add_transform(BM25SearchTransform(always_visible=_ADMIN_TOOL_NAMES))
|
||||
self.server.add_transform(
|
||||
BM25SearchTransform(always_visible=_ADMIN_TOOL_NAMES)
|
||||
)
|
||||
|
||||
def current_config(self) -> BrokerConfig:
|
||||
if self.manager is None:
|
||||
@@ -143,9 +145,7 @@ class TransparentProxyRuntime:
|
||||
) -> dict[str, Any]:
|
||||
tools = await self._list_proxy_tools()
|
||||
if connection_id is not None:
|
||||
tools = [
|
||||
tool for tool in tools if tool["connection_id"] == connection_id
|
||||
]
|
||||
tools = [tool for tool in tools if tool["connection_id"] == connection_id]
|
||||
if query:
|
||||
needle = query.casefold()
|
||||
tools = [
|
||||
|
||||
@@ -298,7 +298,7 @@ class RateChange:
|
||||
@staticmethod
|
||||
def r65(state: CountersContext) -> Rates:
|
||||
c = state.counter
|
||||
assert c["c_80"] >= 65, f"routed wrongly, 80 pity currently at {c["c_80"]}"
|
||||
assert c["c_80"] >= 65, f"routed wrongly, 80 pity currently at {c['c_80']}"
|
||||
n = c["c_80"] - 64
|
||||
rpn = n * 0.05
|
||||
br = state.context["initial_rates"]
|
||||
@@ -401,7 +401,7 @@ def post_roll_router(
|
||||
|
||||
|
||||
@node(name="main")
|
||||
def tick(state: Countdown) -> Countdown: # type: ignore
|
||||
def tick(state: Countdown) -> Countdown:
|
||||
return Countdown(countdown=state.countdown - 1)
|
||||
|
||||
|
||||
@@ -412,23 +412,19 @@ def keep_rolling(state: Countdown) -> NodeReturn[Nothing]:
|
||||
|
||||
# could be @graph.(something combining node and use)...
|
||||
|
||||
gacha.use(
|
||||
init,
|
||||
id="init", # use the node name if not defined?
|
||||
# should i HAVE to declare in / out maps every time?
|
||||
)
|
||||
gacha.use(tick, id="tick")
|
||||
gacha.use(CounterUp.c1, id="counter_up") # 0 base to 1 base probably
|
||||
gacha.use(RateChange.r65, id="rate_up")
|
||||
gacha.use(RateChange.r0, id="rate_same")
|
||||
gacha.use(RateChange.r10, id="r_g10")
|
||||
gacha.use(init)
|
||||
gacha.use(tick, id="tick") # itd use main
|
||||
counter_up = gacha.use(CounterUp.c1, id="counter_up") # 0 base to 1 base probably
|
||||
rate_up = gacha.use(RateChange.r65, id="rate_up")
|
||||
rate_same = gacha.use(RateChange.r0, id="rate_same")
|
||||
r_10 = gacha.use(RateChange.r10, id="r_g10")
|
||||
gacha.use(RateChange.r80, id="r_g80")
|
||||
gacha.use(RateChange.r240, id="r_gs")
|
||||
gacha.use(prep, id="prep")
|
||||
gacha.use(roll, id="roll")
|
||||
gacha.use(CounterUp.c80, id="c_80")
|
||||
prepare_pool = gacha.use(prep)
|
||||
gacha.use(roll)
|
||||
c_80 = gacha.use(CounterUp.c80, id="c_80")
|
||||
gacha.use(CounterUp.c10, id="c_10")
|
||||
gacha.condition(
|
||||
gacha.condition( # condition dont ignore id; you can...
|
||||
id="keep_rolling", check=expr(state("countdown")) > 0
|
||||
) # replaces keep_rolling
|
||||
# Outcome is currently hidden from the docs (there is none), outcome_map is insane, should we have it
|
||||
@@ -440,40 +436,44 @@ gacha.connect("tick", "ok", "counter_up")
|
||||
gacha.use(rate_booster, id="rate_booster")
|
||||
|
||||
gacha.connect("counter_up", "ok", "rate_booster")
|
||||
gacha.connect("rate_booster", "0", "rate_same")
|
||||
gacha.connect("rate_booster", "65", "rate_up")
|
||||
gacha.connect("rate_booster", "0", rate_same)
|
||||
gacha.connect("rate_booster", "65", rate_up)
|
||||
gacha.use(pre_roll_router, id="router")
|
||||
|
||||
gacha.connect("rate_up", "ok", "router")
|
||||
gacha.connect("rate_same", "ok", "router")
|
||||
for outcome, node_id in {
|
||||
"240": "r_gs",
|
||||
"80": "r_g80",
|
||||
"10": "r_g10",
|
||||
"1": "prep",
|
||||
}.items():
|
||||
gacha.connect("router", outcome, node_id)
|
||||
|
||||
|
||||
gacha.connect("prep", "ok", "roll")
|
||||
gacha.connect("r_gs", "ok", "prep")
|
||||
gacha.connect("r_g80", "ok", "prep")
|
||||
gacha.connect("r_g10", "ok", "prep")
|
||||
preroll_routes = gacha.branch(
|
||||
"router",
|
||||
{
|
||||
"240": "r_gs",
|
||||
"80": "r_g80",
|
||||
"10": r_10,
|
||||
"1": "prepare_pool",
|
||||
},
|
||||
)
|
||||
print(preroll_routes)
|
||||
gacha.connect(prepare_pool, "ok", "roll")
|
||||
gacha.connect("r_gs", "ok", "prepare_pool")
|
||||
gacha.connect(preroll_routes["80"], "ok", prepare_pool)
|
||||
gacha.connect("r_g10", "ok", prepare_pool)
|
||||
|
||||
gacha.connect("roll", "ok", "post_roll_router") # missed this! good job
|
||||
gacha.use(post_roll_router, id="post_roll_router")
|
||||
gacha.use(popped, id="reset_avail")
|
||||
gacha.connect("reset_avail", "ok", "c_80")
|
||||
reset_avail = gacha.use(popped, id="reset_avail")
|
||||
gacha.connect(reset_avail, "ok", c_80)
|
||||
|
||||
for outcome, node_id in {
|
||||
"240": "reset_avail",
|
||||
"80": "c_80",
|
||||
"10": "c_10",
|
||||
"1": "keep_rolling",
|
||||
}.items(): # missed this!
|
||||
gacha.connect("post_roll_router", outcome, node_id)
|
||||
# missed this!
|
||||
gacha.branch(
|
||||
"post_roll_router",
|
||||
{
|
||||
"240": "reset_avail",
|
||||
"80": "c_80",
|
||||
"10": "c_10",
|
||||
"1": "keep_rolling",
|
||||
},
|
||||
)
|
||||
|
||||
gacha.connect("c_80", "ok", "keep_rolling")
|
||||
gacha.connect(c_80, "ok", "keep_rolling")
|
||||
gacha.connect("c_10", "ok", "keep_rolling")
|
||||
|
||||
# there is like no general uses for the nodes; idk tho
|
||||
@@ -560,19 +560,24 @@ def execute(graph: WorkflowBuilder, input: Input):
|
||||
c = graph.compile()
|
||||
r = build_registry(*(graph.node_specs.values()))
|
||||
i = input.model_dump()
|
||||
# pprint(i)
|
||||
# pprint(c)
|
||||
# pprint(r)
|
||||
pprint(i)
|
||||
pprint(c)
|
||||
pprint(r)
|
||||
return execute_workflow(c, i, r)
|
||||
|
||||
|
||||
# twice in a row! it took 100+ and a miss tho
|
||||
|
||||
|
||||
def test():
|
||||
d = execute(
|
||||
gacha,
|
||||
build_input(context)(
|
||||
20, rolled_previously=240 - 135, until_5=5, until_6=73, good_stuff=False # lets pretend
|
||||
20,
|
||||
rolled_previously=240 - 135,
|
||||
until_5=5,
|
||||
until_6=73,
|
||||
good_stuff=False, # lets pretend
|
||||
),
|
||||
)
|
||||
assert d.status == RunStatus.COMPLETED, "oops"
|
||||
@@ -580,5 +585,6 @@ def test():
|
||||
assert any(i["name"] in context["pool"]["n_240"] for i in state.storage)
|
||||
pprint(state.storage)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test()
|
||||
test()
|
||||
|
||||
@@ -7,7 +7,7 @@ from pydantic import BaseModel, Field
|
||||
import pytest
|
||||
|
||||
from wf_authoring import WorkflowBuilder, node, state_field
|
||||
from wf_core import RunStatus, WorkflowExecutionError, execute_workflow
|
||||
from wf_core import RunStatus, WorkflowExecutionError
|
||||
|
||||
|
||||
class WorkflowInput(BaseModel):
|
||||
@@ -61,6 +61,12 @@ def auto_bind_node(input: AutoBindInput) -> AutoBindOutput:
|
||||
return AutoBindOutput(text=input.text.upper(), count=input.count + 1)
|
||||
|
||||
|
||||
@node(name="test.branch_router", outcomes=("left", "right"))
|
||||
def branch_router(input: AutoBindInput) -> AutoBindOutput:
|
||||
"""Route to left or right while preserving state shape."""
|
||||
return AutoBindOutput(text=input.text, count=input.count)
|
||||
|
||||
|
||||
def test_builder_accepts_basemodel_classes_for_workflow_schemas() -> None:
|
||||
builder = WorkflowBuilder(
|
||||
name="model_schema_demo",
|
||||
@@ -136,11 +142,8 @@ def test_builder_auto_binds_matching_node_inputs_and_outputs_to_state() -> None:
|
||||
step = builder.use(auto_bind_node, id="update")
|
||||
builder.connect(step, "ok", "__end__")
|
||||
|
||||
workflow = builder.compile()
|
||||
run = execute_workflow(
|
||||
workflow,
|
||||
run = builder.execute(
|
||||
{"text": "hello", "count": 1},
|
||||
builder.registry(),
|
||||
)
|
||||
|
||||
assert step.in_map == {
|
||||
@@ -210,3 +213,74 @@ def test_builder_registry_exports_used_node_specs() -> None:
|
||||
builder.use(auto_bind_node)
|
||||
|
||||
assert set(builder.registry()) == {"test.auto_bind"}
|
||||
|
||||
|
||||
def test_builder_execute_compiles_and_runs_with_used_registry() -> None:
|
||||
builder = WorkflowBuilder(
|
||||
name="execute_demo",
|
||||
input_schema=AutoBindInput,
|
||||
state_schema=AutoBindState,
|
||||
output_schema=AutoBindOutput,
|
||||
)
|
||||
step = builder.use(auto_bind_node)
|
||||
builder.set_entry_point(step)
|
||||
builder.connect(step, "ok", "__end__")
|
||||
|
||||
run = builder.execute({"text": "hello", "count": 1})
|
||||
|
||||
assert run.status == RunStatus.COMPLETED
|
||||
assert run.state["text"] == "HELLO"
|
||||
assert run.state["count"] == 2
|
||||
|
||||
|
||||
def test_builder_branch_connects_existing_steps() -> None:
|
||||
builder = WorkflowBuilder(
|
||||
name="branch_existing_demo",
|
||||
input_schema=AutoBindInput,
|
||||
state_schema=AutoBindState,
|
||||
output_schema=AutoBindOutput,
|
||||
)
|
||||
router = builder.use(branch_router)
|
||||
left = builder.use(auto_bind_node, id="left")
|
||||
right = builder.use(auto_bind_node, id="right")
|
||||
|
||||
builder.branch(router, {"left": left, "right": right})
|
||||
|
||||
assert [(edge.from_, edge.outcome, edge.to) for edge in builder.edges] == [
|
||||
("test_branch_router", "left", "left"),
|
||||
("test_branch_router", "right", "right"),
|
||||
]
|
||||
|
||||
|
||||
def test_builder_branch_can_use_node_specs_as_targets() -> None:
|
||||
builder = WorkflowBuilder(
|
||||
name="branch_specs_demo",
|
||||
input_schema=AutoBindInput,
|
||||
state_schema=AutoBindState,
|
||||
output_schema=AutoBindOutput,
|
||||
)
|
||||
router = builder.use(branch_router)
|
||||
|
||||
targets = builder.branch(router, {"left": auto_bind_node})
|
||||
|
||||
target = targets["left"]
|
||||
assert not isinstance(target, str)
|
||||
assert target.id == "test_auto_bind"
|
||||
assert builder.edges[0].from_ == "test_branch_router"
|
||||
assert builder.edges[0].outcome == "left"
|
||||
assert builder.edges[0].to == "test_auto_bind"
|
||||
|
||||
|
||||
def test_builder_branch_warns_on_empty_branch_map() -> None:
|
||||
builder = WorkflowBuilder(
|
||||
name="branch_empty_demo",
|
||||
input_schema=AutoBindInput,
|
||||
state_schema=AutoBindState,
|
||||
output_schema=AutoBindOutput,
|
||||
)
|
||||
router = builder.use(branch_router)
|
||||
|
||||
with pytest.warns(UserWarning, match="no branches"):
|
||||
targets = builder.branch(router, {})
|
||||
|
||||
assert targets == {}
|
||||
|
||||
@@ -265,9 +265,10 @@ def test_transparent_proxy_proxy_tool_listing_supports_filters_and_cursor() -> N
|
||||
)
|
||||
second_page = _structured(second_page_result)
|
||||
assert len(second_page["tools"]) == 1
|
||||
assert second_page["tools"][0]["proxy_name"] != first_page["tools"][0][
|
||||
"proxy_name"
|
||||
]
|
||||
assert (
|
||||
second_page["tools"][0]["proxy_name"]
|
||||
!= first_page["tools"][0]["proxy_name"]
|
||||
)
|
||||
|
||||
filtered_result = await client.call_tool(
|
||||
"wf.mcp_list_proxy_tools",
|
||||
|
||||
Reference in New Issue
Block a user