diff --git a/burr/core/__init__.py b/burr/core/__init__.py index ddc1f53e6..9de9e9100 100644 --- a/burr/core/__init__.py +++ b/burr/core/__init__.py @@ -30,6 +30,7 @@ Application, ApplicationBuilder, ApplicationContext, + ApplicationExecutionLimitError, ApplicationGraph, ) from burr.core.graph import Graph, GraphBuilder @@ -42,6 +43,7 @@ "ApplicationBuilder", "ApplicationGraph", "ApplicationContext", + "ApplicationExecutionLimitError", "Condition", "default", "expr", diff --git a/burr/core/application.py b/burr/core/application.py index 25bce4a10..202f5498d 100644 --- a/burr/core/application.py +++ b/burr/core/application.py @@ -89,6 +89,19 @@ StateTypeToSet = TypeVar("StateTypeToSet") +class ApplicationExecutionLimitError(RuntimeError): + """Raised when an application execution call exhausts its step budget.""" + + def __init__(self, max_steps: int, last_action: Action, state: State): + self.max_steps = max_steps + self.last_action = last_action + self.state = state + super().__init__( + f"Application execution reached max_steps={max_steps} after action " + f"'{last_action.name}' while another action was still available." + ) + + def _validate_result(result: Any, name: str, schema: ActionSchema = DEFAULT_SCHEMA) -> None: # TODO -- split out the action schema into input/output schema types # Currently they're tied together, but this doesn't make as much sense for single-step actions @@ -1262,6 +1275,7 @@ def iterate( halt_before: list[str] = None, halt_after: list[str] = None, inputs: Optional[Dict[str, Any]] = None, + max_steps: Optional[int] = None, ) -> Generator[ Tuple[Action, dict, State[ApplicationStateType]], None, @@ -1280,6 +1294,7 @@ def iterate( :param halt_after: The list of actions/tags to halt after execution of. It will halt after the execution of the first one it sees. :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world. Note that this is only used for the first iteration -- subsequent iterations will not use this. + :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution. :return: Each iteration returns the result of running `step`. This generator also returns a tuple of [action, result, current state] """ @@ -1290,15 +1305,23 @@ def iterate( inputs, ) self._validate_halt_conditions(halt_before, halt_after) + if max_steps is not None and ( + not isinstance(max_steps, int) or isinstance(max_steps, bool) or max_steps <= 0 + ): + raise ValueError("max_steps must be a positive integer or None") result = None prior_action: Optional[Action] = None + steps_executed = 0 while self.has_next_action(): # self.step will only return None if there is no next action, so we can rely on tuple unpacking prior_action, result, state = self.step(inputs=inputs) + steps_executed += 1 yield prior_action, result, state if self._should_halt_iterate(halt_before, halt_after, prior_action): break + if max_steps is not None and steps_executed >= max_steps and self.has_next_action(): + raise ApplicationExecutionLimitError(max_steps, prior_action, state) return self._return_value_iterate(halt_before, halt_after, prior_action, result) @_call_execute_method_pre_post(ExecuteMethod.aiterate) @@ -1308,6 +1331,7 @@ async def aiterate( halt_before: list[str] = None, halt_after: list[str] = None, inputs: Optional[Dict[str, Any]] = None, + max_steps: Optional[int] = None, ) -> AsyncGenerator[Tuple[Action, dict, State[ApplicationStateType]], None]: """Returns a generator that calls step() in a row, enabling you to see the state of the system as it updates. This is the asynchronous version so it has no capability of t @@ -1318,6 +1342,7 @@ async def aiterate( :param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one. :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world. Note that this is only used for the first iteration -- subsequent iterations will not use this. + :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution. :return: Each iteration returns the result of running `step`. This returns nothing -- it's an async generator which is not allowed to have a return value. """ @@ -1325,12 +1350,20 @@ async def aiterate( halt_before, halt_after, inputs ) self._validate_halt_conditions(halt_before, halt_after) + if max_steps is not None and ( + not isinstance(max_steps, int) or isinstance(max_steps, bool) or max_steps <= 0 + ): + raise ValueError("max_steps must be a positive integer or None") + steps_executed = 0 while self.has_next_action(): # self.step will only return None if there is no next action, so we can rely on tuple unpacking prior_action, result, state = await self.astep(inputs=inputs) + steps_executed += 1 yield prior_action, result, state if self._should_halt_iterate(halt_before, halt_after, prior_action): break + if max_steps is not None and steps_executed >= max_steps and self.has_next_action(): + raise ApplicationExecutionLimitError(max_steps, prior_action, state) @_call_execute_method_pre_post(ExecuteMethod.run) def run( @@ -1339,6 +1372,7 @@ def run( halt_before: list[str] = None, halt_after: list[str] = None, inputs: Optional[Dict[str, Any]] = None, + max_steps: Optional[int] = None, ) -> Tuple[Action, Optional[dict], State[ApplicationStateType]]: """Runs your application through until completion. Does not give access to the state along the way -- if you want that, use iterate(). @@ -1349,10 +1383,16 @@ def run( :param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one. :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world. Note that this is only used for the first iteration -- subsequent iterations will not use this. + :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution. :return: The final state, and the results of running the actions in the order that they were specified. """ self.validate_correct_async_use() - gen = self.iterate(halt_before=halt_before, halt_after=halt_after, inputs=inputs) + gen = self.iterate( + halt_before=halt_before, + halt_after=halt_after, + inputs=inputs, + max_steps=max_steps, + ) while True: try: next(gen) @@ -1367,6 +1407,7 @@ async def arun( halt_before: list[str] = None, halt_after: list[str] = None, inputs: Optional[Dict[str, Any]] = None, + max_steps: Optional[int] = None, ) -> Tuple[Action, Optional[dict], State[ApplicationStateType]]: """Runs your application through until completion, using async. Does not give access to the state along the way -- if you want that, use iterate(). @@ -1376,6 +1417,7 @@ async def arun( :param halt_before: The list of actions/tags to halt before execution of. It will halt on the first one. :param halt_after: The list of actions/tags to halt after execution of. It will halt on the first one. :param inputs: Inputs to the action -- this is if this action requires an input that is passed in from the outside world + :param max_steps: Maximum actions to execute during this call. ``None`` preserves unbounded execution. :return: The final state, and the results of running the actions in the order that they were specified. """ @@ -1386,7 +1428,10 @@ async def arun( ) self._validate_halt_conditions(halt_before, halt_after) async for prior_action, result, state in self.aiterate( - halt_before=halt_before, halt_after=halt_after, inputs=inputs + halt_before=halt_before, + halt_after=halt_after, + inputs=inputs, + max_steps=max_steps, ): pass return self._return_value_iterate(halt_before, halt_after, prior_action, result) diff --git a/docs/concepts/state-machine.rst b/docs/concepts/state-machine.rst index 2fea48aea..0b418bdc4 100644 --- a/docs/concepts/state-machine.rst +++ b/docs/concepts/state-machine.rst @@ -116,6 +116,28 @@ In the async context, this does not return anything (asynchronous generators are If you want it to (attempt to) run forever, you can pass empty lists to ``halt_after`` and ``halt_before``. +Limiting execution steps +~~~~~~~~~~~~~~~~~~~~~~~~ + +Agent and tool-use graphs often contain intentional loops. Pass ``max_steps`` to +``run``/``arun`` or ``iterate``/``aiterate`` to bound the number of actions a +single call may execute: + +.. code-block:: python + + from burr.core import ApplicationExecutionLimitError + + try: + action, result, state = application.run(max_steps=20) + except ApplicationExecutionLimitError as exc: + # The application remains at the state produced by the last allowed step. + save_for_review(exc.state) + +The exception also exposes ``max_steps`` and ``last_action``. A halt condition +or the natural end of the graph takes precedence when it is reached on the last +allowed step. The budget applies only to the current execution call, so the +application can be inspected and resumed with a new call. + .. note:: You can add inputs to ``iterate``/``aiterate`` by passing in a dictionary of inputs through the ``inputs`` parameter. This will only apply to the first action. Actions that are not the first but require inputs are considered undefined behavior. diff --git a/docs/reference/application.rst b/docs/reference/application.rst index d1a03d82e..e3a5920f3 100644 --- a/docs/reference/application.rst +++ b/docs/reference/application.rst @@ -41,6 +41,9 @@ and not the ``Application`` class directly. .. autoclass:: burr.core.application.ApplicationContext :members: +.. autoclass:: burr.core.application.ApplicationExecutionLimitError + :members: + ========== Graph APIs ========== diff --git a/tests/core/test_application.py b/tests/core/test_application.py index 9313cefc9..3a8c1f912 100644 --- a/tests/core/test_application.py +++ b/tests/core/test_application.py @@ -25,7 +25,7 @@ import pytest -from burr.core import State +from burr.core import ApplicationExecutionLimitError, State from burr.core.action import ( DEFAULT_SCHEMA, Action, @@ -2171,6 +2171,184 @@ async def test_aiterate(): assert app.sequence_id == 11 +def test_run_stops_before_exceeding_max_steps(): + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph( + actions=[counter_action], + transitions=[Transition(counter_action, counter_action, default)], + ), + ) + + with pytest.raises(ApplicationExecutionLimitError) as exc_info: + app.run(max_steps=3) + + assert exc_info.value.max_steps == 3 + assert exc_info.value.last_action.name == "counter" + assert exc_info.value.state["count"] == 3 + assert app.state["count"] == 3 + + +def test_iterate_raises_only_when_requesting_a_step_beyond_the_budget(): + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph( + actions=[counter_action], + transitions=[Transition(counter_action, counter_action, default)], + ), + ) + iterator = app.iterate(max_steps=1) + + action_, result, state = next(iterator) + + assert action_.name == "counter" + assert result == {"count": 1} + assert state["count"] == 1 + with pytest.raises(ApplicationExecutionLimitError): + next(iterator) + + +def test_run_allows_natural_completion_at_max_steps(): + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + action_, result, state = app.run(max_steps=1) + + assert action_.name == "counter" + assert result == {"count": 1} + assert state["count"] == 1 + + +def test_run_can_resume_after_reaching_max_steps(): + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph( + actions=[counter_action], + transitions=[Transition(counter_action, counter_action, default)], + ), + ) + with pytest.raises(ApplicationExecutionLimitError): + app.run(max_steps=1) + + action_, result, state = app.run(halt_after=["counter"], max_steps=1) + + assert action_.name == "counter" + assert result == {"count": 2} + assert state["count"] == 2 + + +def test_run_halt_condition_takes_precedence_over_max_steps(): + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph( + actions=[counter_action], + transitions=[Transition(counter_action, counter_action, default)], + ), + ) + + action_, result, state = app.run(halt_after=["counter"], max_steps=1) + + assert action_.name == "counter" + assert result == {"count": 1} + assert state["count"] == 1 + + +@pytest.mark.parametrize("max_steps", [0, -1, 1.5, "3", True]) +def test_run_rejects_invalid_max_steps(max_steps): + counter_action = base_counter_action.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph(actions=[counter_action], transitions=[]), + ) + + with pytest.raises(ValueError, match="max_steps must be a positive integer"): + app.run(max_steps=max_steps) + + assert app.sequence_id == 0 + + +async def test_arun_stops_before_exceeding_max_steps(): + counter_action = base_counter_action_async.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph( + actions=[counter_action], + transitions=[Transition(counter_action, counter_action, default)], + ), + ) + + with pytest.raises(ApplicationExecutionLimitError) as exc_info: + await app.arun(max_steps=2) + + assert exc_info.value.max_steps == 2 + assert exc_info.value.last_action.name == "counter" + assert exc_info.value.state["count"] == 2 + + +async def test_aiterate_raises_after_yield_and_async_halt_can_resume(): + counter_action = base_counter_action_async.with_name("counter") + app = Application( + state=State({}), + entrypoint="counter", + partition_key="test", + uid="test-123", + sequence_id=0, + graph=Graph( + actions=[counter_action], + transitions=[Transition(counter_action, counter_action, default)], + ), + ) + iterator = app.aiterate(max_steps=1) + + action_, result, state = await iterator.__anext__() + + assert action_.name == "counter" + assert result == {"count": 1} + assert state["count"] == 1 + with pytest.raises(ApplicationExecutionLimitError): + await iterator.__anext__() + + action_, result, state = await app.arun(halt_after=["counter"], max_steps=1) + assert action_.name == "counter" + assert result == {"count": 2} + assert state["count"] == 2 + + async def test_aiterate_halt_before(): result_action = Result("count").with_name("result") counter_action = base_counter_action_async.with_name("counter")