and more good stuff, fmt

This commit is contained in:
lda
2026-05-07 05:03:35 +07:00 Verified
parent 8fea1769fe
commit 7a13e05a3d
7 changed files with 193 additions and 67 deletions
+48 -2
View File
@@ -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(
+4 -6
View File
@@ -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
View File
@@ -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
+4 -4
View File
@@ -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 = [
+51 -45
View File
@@ -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()
+79 -5
View File
@@ -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 == {}
+4 -3
View File
@@ -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",