diff --git a/burr/lifecycle/__init__.py b/burr/lifecycle/__init__.py index 4ae24073a..fa3a5f90c 100644 --- a/burr/lifecycle/__init__.py +++ b/burr/lifecycle/__init__.py @@ -29,7 +29,12 @@ PreRunStepHookAsync, PreStartSpanHook, ) -from burr.lifecycle.default import StateAndResultsFullLogger +from burr.lifecycle.default import ( + ExecutionRecord, + InMemoryExecutionRecorder, + StateAndResultsFullLogger, + StateChange, +) __all__ = [ "PreRunStepHook", @@ -42,6 +47,9 @@ "PostApplicationExecuteCallHookAsync", "LifecycleAdapter", "StateAndResultsFullLogger", + "InMemoryExecutionRecorder", + "ExecutionRecord", + "StateChange", "PostApplicationCreateHook", "PostEndSpanHook", "PreStartSpanHook", diff --git a/burr/lifecycle/default.py b/burr/lifecycle/default.py index 8bc234a26..e98b489c0 100644 --- a/burr/lifecycle/default.py +++ b/burr/lifecycle/default.py @@ -15,10 +15,12 @@ # specific language governing permissions and limitations # under the License. +import copy +import dataclasses import datetime import json import time -from typing import TYPE_CHECKING, Any, Callable, Literal, Optional +from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional if TYPE_CHECKING: from burr.core import State, Action @@ -30,6 +32,139 @@ def safe_json(obj: Any) -> str: return json.dumps(obj, default=str) +def _values_equal(left: Any, right: Any) -> bool: + if left is right: + return True + try: + comparison = left == right + except Exception: + return False + return comparison if isinstance(comparison, bool) else False + + +def _copy_for_recording(value: Any) -> Any: + try: + return copy.deepcopy(value) + except Exception: + return value + + +def _copy_mapping_for_recording(values: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + if values is None: + return None + return {key: _copy_for_recording(value) for key, value in values.items()} + + +@dataclasses.dataclass(frozen=True) +class StateChange: + """A single state field change captured around an action execution.""" + + key: str + before_exists: bool + before: Any + after_exists: bool + after: Any + + +@dataclasses.dataclass(frozen=True) +class ExecutionRecord: + """An immutable record of one Burr action execution.""" + + app_id: str + partition_key: Optional[str] + sequence_id: int + action: str + inputs: Dict[str, Any] + result: Optional[Dict[str, Any]] + exception: Optional[Exception] + state_changes: tuple[StateChange, ...] + + +class InMemoryExecutionRecorder(PreRunStepHook, PostRunStepHook): + """Capture action executions in memory for tests and local debugging. + + The recorder stores inputs, results, exceptions, and business-state changes. + Burr's private ``__`` state fields are excluded from the change list. + """ + + def __init__(self): + self._records: list[ExecutionRecord] = [] + self._pending: dict[tuple[str, Optional[str], int], tuple[dict, dict]] = {} + + @property + def records(self) -> tuple[ExecutionRecord, ...]: + """Return an immutable snapshot of captured execution records.""" + return tuple(self._records) + + def clear(self) -> None: + """Remove all completed and pending records.""" + self._records.clear() + self._pending.clear() + + def pre_run_step( + self, + *, + app_id: str, + partition_key: Optional[str], + sequence_id: int, + state: "State", + inputs: Dict[str, Any], + **future_kwargs: Any, + ) -> None: + key = (app_id, partition_key, sequence_id) + self._pending[key] = ( + _copy_mapping_for_recording(state.get_all()), + _copy_mapping_for_recording(inputs), + ) + + def post_run_step( + self, + *, + app_id: str, + partition_key: Optional[str], + sequence_id: int, + state: "State", + action: "Action", + result: Optional[Dict[str, Any]], + exception: Optional[Exception], + **future_kwargs: Any, + ) -> None: + key = (app_id, partition_key, sequence_id) + before, inputs = self._pending.pop(key) + after = state.get_all() + changed_keys = sorted( + key + for key in before.keys() | after.keys() + if not key.startswith("__") + and ( + (key in before) != (key in after) + or not _values_equal(before.get(key), after.get(key)) + ) + ) + changes = tuple( + StateChange( + key=state_key, + before_exists=state_key in before, + before=_copy_for_recording(before.get(state_key)), + after_exists=state_key in after, + after=_copy_for_recording(after.get(state_key)), + ) + for state_key in changed_keys + ) + self._records.append( + ExecutionRecord( + app_id=app_id, + partition_key=partition_key, + sequence_id=sequence_id, + action=action.name, + inputs=inputs, + result=_copy_mapping_for_recording(result), + exception=exception, + state_changes=changes, + ) + ) + + class StateAndResultsFullLogger(PostRunStepHook, PreRunStepHook): """Logs the state and results of the action in a jsonl file.""" diff --git a/docs/concepts/hooks.rst b/docs/concepts/hooks.rst index dbd62f12f..fdbccf942 100644 --- a/docs/concepts/hooks.rst +++ b/docs/concepts/hooks.rst @@ -74,6 +74,32 @@ To include this in the application, you pass it into the :py:meth:`with_hooks State: + return state.update(count=state.get("count", 0) + amount) + + +@action(reads=[], writes=[]) +def fail(state: State) -> State: + raise RuntimeError("tool failed") + + +@action(reads=[], writes=["nullable"]) +def add_nullable(state: State) -> State: + return state.update(nullable=None) + + +@action(reads=["count"], writes=["count"]) +async def increment_async(state: State) -> State: + return state.update(count=state.get("count", 0) + 1) + + +def test_in_memory_execution_recorder_captures_step_changes(): + recorder = InMemoryExecutionRecorder() + app = ( + ApplicationBuilder() + .with_actions(increment) + .with_entrypoint("increment") + .with_state(count=1) + .with_identifiers(app_id="agent-run", partition_key="user-1") + .with_hooks(recorder) + .build() + ) + + app.step(inputs={"amount": 2}) + + assert len(recorder.records) == 1 + record = recorder.records[0] + assert record.app_id == "agent-run" + assert record.partition_key == "user-1" + assert record.sequence_id == 0 + assert record.action == "increment" + assert record.inputs == {"amount": 2} + assert record.result == {} + assert record.exception is None + assert record.state_changes == ( + StateChange( + key="count", + before_exists=True, + before=1, + after_exists=True, + after=3, + ), + ) + + +def test_in_memory_execution_recorder_captures_failed_step(): + recorder = InMemoryExecutionRecorder() + app = ( + ApplicationBuilder().with_actions(fail).with_entrypoint("fail").with_hooks(recorder).build() + ) + + with pytest.raises(RuntimeError, match="tool failed"): + app.step() + + assert len(recorder.records) == 1 + record = recorder.records[0] + assert record.action == "fail" + assert record.result is None + assert isinstance(record.exception, RuntimeError) + assert record.state_changes == () + + +def test_in_memory_execution_recorder_clear_removes_records(): + recorder = InMemoryExecutionRecorder() + app = ( + ApplicationBuilder() + .with_actions(increment) + .with_entrypoint("increment") + .with_hooks(recorder) + .build() + ) + app.step() + + recorder.clear() + + assert recorder.records == () + + +def test_in_memory_execution_recorder_distinguishes_missing_from_none(): + recorder = InMemoryExecutionRecorder() + app = ( + ApplicationBuilder() + .with_actions(add_nullable) + .with_entrypoint("add_nullable") + .with_hooks(recorder) + .build() + ) + + app.step() + + assert recorder.records[0].state_changes == ( + StateChange( + key="nullable", + before_exists=False, + before=None, + after_exists=True, + after=None, + ), + ) + + +async def test_in_memory_execution_recorder_captures_async_steps(): + recorder = InMemoryExecutionRecorder() + app = ( + ApplicationBuilder() + .with_actions(increment_async) + .with_entrypoint("increment_async") + .with_hooks(recorder) + .with_state(count=0) + .build() + ) + + await app.astep() + + assert recorder.records[0].action == "increment_async" + assert recorder.records[0].state_changes[0].after == 1 + + +def test_in_memory_execution_recorder_handles_non_boolean_equality(): + class VectorLike: + def __deepcopy__(self, memo): + return self + + def __eq__(self, other): + return self + + def __bool__(self): + raise ValueError("truth value is ambiguous") + + recorder = InMemoryExecutionRecorder() + app = ( + ApplicationBuilder() + .with_actions(increment) + .with_entrypoint("increment") + .with_state(count=0, embeddings=VectorLike()) + .with_hooks(recorder) + .build() + ) + + app.step() + + assert [change.key for change in recorder.records[0].state_changes] == ["count"] + + +def test_in_memory_execution_recorder_does_not_block_uncopyable_state(): + class Uncopyable: + def __deepcopy__(self, memo): + raise TypeError("cannot copy runtime handle") + + recorder = InMemoryExecutionRecorder() + app = ( + ApplicationBuilder() + .with_actions(increment) + .with_entrypoint("increment") + .with_state(count=0, handle=Uncopyable()) + .with_hooks(recorder) + .build() + ) + + action_, result, state = app.step() + + assert action_.name == "increment" + assert result == {} + assert state["count"] == 1 + assert [change.key for change in recorder.records[0].state_changes] == ["count"]