added Select to define transitions

burr/core/__init__.py

@@ -1,4 +1,4 @@
-from burr.core.action import Action, Condition, Result, action, default, expr, when
+from burr.core.action import Action, Condition, Result, Select, action, default, expr, when
 from burr.core.application import (
  Application,
  ApplicationBuilder,
@@ -18,6 +18,7 @@
  "default",
  "expr",
  "Result",
+ "Select",
  "State",
  "when",
 ]

burr/core/action.py

@@ -17,6 +17,7 @@
  List,
  Optional,
  Protocol,
+ Sequence,
  Tuple,
  TypeVar,
  Union,
@@ -378,6 +379,50 @@ def __invert__(self):
 # exists = Condition.exists
 
 
+# TODO type `resolver` to prevent user-facing type-mismatch
+# e.g., a user provided `def foo(state: State, actions: list)`
+# would be too restrictive for a `Sequence` type
+class Select(Function):
+ def __init__(
+ self,
+ keys: List[str],
+ resolver: Callable[[State, Sequence[Action]], str],
+ name: str = None,
+ ):
+ self._keys = keys
+ self._resolver = resolver
+ self._name = name
+ # TODO add a `default` kwarg;
+ # could an Action, action_name: str, or action_idx: int
+ # `default` value could be returned if `_resolver` returns None
+
+ @property
+ def name(self) -> str:
+ return self._name
+
+ @property
+ def reads(self) -> list[str]:
+ return self._keys
+
+ @property
+ def resolver(self) -> Callable[[State, Sequence[Action]], str]:
+ return self._resolver
+
+ def __repr__(self) -> str:
+ return f"select: {self._name}"
+
+ def _validate(self, state: State):
+ missing_keys = set(self._keys) - set(state.keys())
+ if missing_keys:
+ raise ValueError(
+ f"Missing keys in state required by condition: {self} {', '.join(missing_keys)}"
+ )
+
+ def run(self, state: State, possible_actions: Sequence[Action]) -> str:
+ self._validate(state)
+ return self._resolver(state, possible_actions)
+
+
 class Result(Action):
  def __init__(self, *fields: str):
  """Represents a result action. This is purely a convenience class to

burr/core/application.py

@@ -16,6 +16,7 @@
  List,
  Literal,
  Optional,
+ Sequence,
  Set,
  Tuple,
  TypeVar,
@@ -33,6 +34,7 @@
  Condition,
  Function,
  Reducer,
+ Select,
  SingleStepAction,
  SingleStepStreamingAction,
  StreamingAction,
@@ -2005,7 +2007,8 @@ def with_actions(
  def with_transitions(
  self,
  *transitions: Union[
- Tuple[Union[str, list[str]], str], Tuple[Union[str, list[str]], str, Condition]
+ Tuple[Union[str, Sequence[str]], Union[str, Sequence[str]]],
+ Tuple[Union[str, Sequence[str]], Union[str, Sequence[str]], Union[Condition, Select]],
  ],
  ) -> "ApplicationBuilder":
  """Adds transitions to the application. Transitions are specified as tuples of either:

burr/core/graph.py

@@ -3,10 +3,10 @@
 import inspect
 import logging
 import pathlib
-from typing import Any, Callable, List, Literal, Optional, Set, Tuple, Union
+from typing import Any, Callable, List, Literal, Optional, Sequence, Set, Tuple, Union
 
 from burr import telemetry
-from burr.core.action import Action, Condition, create_action, default
+from burr.core.action import Action, Condition, Select, create_action, default
 from burr.core.state import State
 from burr.core.validation import BASE_ERROR_MESSAGE, assert_set
 
@@ -118,6 +118,13 @@ def get_next_node(
  return self._action_map[entrypoint]
  possibilities = self._adjacency_map[prior_step]
  for next_action, condition in possibilities:
+ # When `Select` is used, all possibilities have the same `condition` attached.
+ # Hitting a `Select` will necessarily exit the for loop
+ if isinstance(condition, Select):
+ possible_actions = [self._action_map[p[0]] for p in possibilities]
+ selected_action = condition.run(state, possible_actions)
+ return self._action_map[selected_action]
+
  if condition.run(state)[Condition.KEY]:
  return self._action_map[next_action]
  return None
@@ -235,7 +242,7 @@ class GraphBuilder:
 
  def __init__(self):
  """Initializes the graph builder."""
- self.transitions: Optional[List[Tuple[str, str, Condition]]] = None
+ self.transitions: Optional[List[Tuple[str, str, Union[Condition, Select]]]] = None
  self.actions: Optional[List[Action]] = None
 
  def with_actions(
@@ -269,7 +276,8 @@ def with_actions(
  def with_transitions(
  self,
  *transitions: Union[
- Tuple[Union[str, list[str]], str], Tuple[Union[str, list[str]], str, Condition]
+ Tuple[Union[str, Sequence[str]], Union[str, Sequence[str]]],
+ Tuple[Union[str, Sequence[str]], Union[str, Sequence[str]], Union[Condition, Select]],
  ],
  ) -> "GraphBuilder":
  """Adds transitions to the graph. Transitions are specified as tuples of either:
@@ -291,14 +299,25 @@ def with_transitions(
  condition = conditions[0]
  else:
  condition = default
- if not isinstance(from_, list):
+ # check required because issubclass(str, Sequence) == True
+ if isinstance(from_, Sequence) and not isinstance(from_, str):
+ from_ = [*from_]
+ else:
  from_ = [from_]
- for action in from_:
- if not isinstance(action, str):
- raise ValueError(f"Transition source must be a string, not {action}")
- if not isinstance(to_, str):
- raise ValueError(f"Transition target must be a string, not {to_}")
- self.transitions.append((action, to_, condition))
+ if isinstance(to_, Sequence) and not isinstance(to_, str):
+ if not isinstance(condition, Select):
+ raise ValueError(
+ "Transition with multiple targets require a `Select` condition."
+ )
+ else:
+ to_ = [to_]
+ for source in from_:
+ for target in to_:
+ if not isinstance(source, str):
+ raise ValueError(f"Transition source must be a string, not {source}")
+ if not isinstance(target, str):
+ raise ValueError(f"Transition target must be a string, not {to_}")
+ self.transitions.append((source, target, condition))
  return self
 
  def build(self) -> Graph:

tests/core/test_action.py

@@ -12,6 +12,7 @@
  Function,
  Input,
  Result,
+ Select,
  SingleStepAction,
  SingleStepStreamingAction,
  StreamingAction,
@@ -200,6 +201,39 @@ def test_condition_lmda():
 # assert cond.run(State({"foo" : "bar"})) == {Condition.KEY: False}
 
 
+def test_select_constant():
+ select = Select([], resolver=lambda *args: "foo")
+ selected_action = select.run(State(), [])
+
+ assert selected_action == "foo"
+
+
+def test_select_determistic():
+ @action(reads=[], writes=[])
+ def bar(state):
+ return state
+
+ @action(reads=[], writes=[])
+ def baz(state):
+ return state
+
+ def length_resolver(state: State, actions: list[Action]) -> str:
+ foo = state["foo"]
+ action_idx = len(foo) % len(actions)
+ return actions[action_idx].name
+
+ foo1 = "len=3" # % 2 = 1
+ foo2 = "len_is_8" # % 2 = 0
+ actions = [create_action(bar, "bar"), create_action(baz, "baz")]
+ select = Select(["foo"], resolver=length_resolver)
+
+ selected_1 = select.run(State({"foo": foo1}), possible_actions=actions)
+ assert selected_1 == actions[len(foo1) % len(actions)].name
+
+ selected_2 = select.run(State({"foo": foo2}), possible_actions=actions)
+ assert selected_2 == actions[len(foo2) % len(actions)].name
+
+
 def test_result():
  result = Result("foo", "bar")
  assert result.run(State({"foo": "baz", "bar": "qux", "baz": "quux"})) == {