Adds ability to fork application via initialize_from

burr/core/application.py

@@ -1204,6 +1204,9 @@ def __init__(self):
         self.initializer = None
         self.use_entrypoint_from_save_state: Optional[bool] = None
         self.default_state: Optional[dict] = None
+        self.fork_from_app_id: Optional[str] = None
+        self.fork_from_partition_key: Optional[str] = None
+        self.fork_from_sequence_id: Optional[int] = None
 
     def with_identifiers(
         self, app_id: str = None, partition_key: str = None, sequence_id: int = None
@@ -1386,15 +1389,24 @@ def initialize_from(
         resume_at_next_action: bool,
         default_state: dict,
         default_entrypoint: str,
+        fork_from_app_id: str = None,
+        fork_from_partition_key: str = None,
+        fork_from_sequence_id: int = None,
     ) -> "ApplicationBuilder":
         """Initializes the application we will build from some prior state object.
-        Note that you can *either* call this or use `with_state` and `with_entrypoint` -- this also assigns application ID,
-        partition key, and sequence ID.
+
+        Note (1) that you can *either* call this or use `with_state` and `with_entrypoint`.
+
+        Note (2) if you want to continue a prior application and don't want to fork it into a new application ID,
+        the values in `.with_identifiers()` will be used to query for prior state.
 
         :param initializer: The persister object to use for initialization. Likely the same one called with ``with_state_persister``.
         :param resume_at_next_action: Whether to resume at the next action, or default to the ``default_entrypoint``
         :param default_state: The default state to use if it does not exist. This is a dictionary.
         :param default_entrypoint: The default entry point to use if it does not exist or you elect not to resume_at_next_action.
+        :param fork_from_app_id: The app ID to fork from, not to be confused with the current app_id that is set with `.with_identifiers()`. This is used to fork from a prior application run.
+        :param fork_from_partition_key: The partition key to fork from a prior application. Optional. `fork_from_app_id` required.
+        :param fork_from_sequence_id: The sequence ID to fork from a prior application run. Optional, defaults to latest. `fork_from_app_id` required.
         :return: The application builder for future chaining.
         """
         if self.start is not None or self.state is not None:
@@ -1403,10 +1415,19 @@ def initialize_from(
                 + "Cannot call initialize_from if you have already set state or an entrypoint! "
                 "You can either use the initializer *or* set the state and entrypoint manually."
             )
+        if not fork_from_app_id and (fork_from_partition_key or fork_from_sequence_id):
+            raise ValueError(
+                ERROR_MESSAGE
+                + "If you set fork_from_partition_key or fork_from_sequence_id, you must also set fork_from_app_id. "
+                "See .initialize_from() documentation."
+            )
         self.initializer = initializer
         self.resume_at_next_action = resume_at_next_action
         self.default_state = default_state
         self.start = default_entrypoint
+        self.fork_from_app_id = fork_from_app_id
+        self.fork_from_partition_key = fork_from_partition_key
+        self.fork_from_sequence_id = fork_from_sequence_id
         return self
 
     def with_state_persister(
@@ -1441,9 +1462,32 @@ def _load_from_persister(self):
          - maybe self.start
 
         """
+        if self.fork_from_app_id is not None:
+            if self.app_id == self.fork_from_app_id:
+                raise ValueError(
+                    ERROR_MESSAGE + "Cannot fork and save to the same app_id. "
+                    "Please update the app_id passed in via with_identifiers(), "
+                    "or don't pass in a fork_from_app_id value to `initialize_from()`."
+                )
+            _partition_key = self.fork_from_partition_key
+            _app_id = self.fork_from_app_id
+            _sequence_id = self.fork_from_sequence_id
+        else:
+            # only use the with_identifier values if we're not forking from a previous app
+            _partition_key = self.partition_key
+            _app_id = self.app_id
+            _sequence_id = self.sequence_id
         # load state from persister
-        load_result = self.initializer.load(self.partition_key, self.app_id, self.sequence_id)
+        load_result = self.initializer.load(_partition_key, _app_id, _sequence_id)
         if load_result is None:
+            if self.fork_from_app_id is not None:
+                logger.warning(
+                    f"{self.initializer.__class__.__name__} returned None while trying to fork from: "
+                    f"partition_key:{_partition_key}, app_id:{_app_id}, "
+                    f"sequence_id:{_sequence_id}. "
+                    "You explicitly requested to fork from a prior application run, but it does not exist. "
+                    "Defaulting to state defaults instead."
+                )
             # there was nothing to load -- use default state
             self.state = self.state.update(**self.default_state)
             self.sequence_id = None  # has to start at None
@@ -1452,11 +1496,12 @@ def _load_from_persister(self):
                 raise ValueError(
                     ERROR_MESSAGE
                     + f"Error: {self.initializer.__class__.__name__} returned {load_result} for "
-                    f"partition_key:{self.partition_key}, app_id:{self.app_id}, "
-                    f"sequence_id:{self.sequence_id}, "
-                    f"but state was None! This is not allowed. Please return None in this case, or double "
-                    f"check that persisted state can never be a None value."
+                    f"partition_key:{_partition_key}, app_id:{_app_id}, "
+                    f"sequence_id:{_sequence_id}, "
+                    "but value for state was None! This is not allowed. Please return just None in this case, "
+                    "or double check that persisted state can never be a None value."
                 )
+            # TODO: capture parent app ID relationship & wire it through
             # there was something
             last_position = load_result["position"]
             self.state = load_result["state"]

docs/concepts/state-persistence.rst

@@ -27,21 +27,41 @@ Burr `applications` are, by defult, keyed on two entities:
 In the case of a chatbot, the ``app_uid`` could be a uuid, and the ``partition_key`` could be the user's name.
 Note that ``partition_key`` can be `None` if this is not relevant. A UUID is always generated for the ``app_uid`` if not provided.
 
+You set these values using the :py:meth:`with_identifiers() <burr.core.application.ApplicationBuilder.with_identifiers>` method.
+
 Initializing state
 ------------------
 
 To initialize state from a database, you can employ the :py:meth:`initialize_from <burr.core.application.ApplicationBuilder.initialize_from>` method
-in the :py:class:`ApplicationBuilder <burr.core.application.ApplicationBuilder>`.
+in the :py:class:`ApplicationBuilder <burr.core.application.ApplicationBuilder>`. Important API note: unless `fork_from_app_id` is specified
+below, the values used to query for state are taken from the values provided to :py:meth:`with_identifiers() <burr.core.application.ApplicationBuilder.with_identifiers>`.
 
 This action takes in an initializer (an implementation of :py:class:`StateInitializer <burr.core.persistence.BaseStateLoader>`) a well as:
 
 - ``resume_at_next_action`` -- a boolean that says whether to start where you left off, or go back to the ``default_entrypoint``.
 - ``default_entrypoint`` -- the entry point to start at if ``resume_at_next_action`` is False, or no state is found
 - ``default_state`` -- the default state to use if no state is found
+- ``fork_from_app_id`` -- Optional. A prior app_id to fork state from. This is useful if you want to start from a previous application's state.
+- ``fork_from_partition_key`` -- Optional. The partition key to fork from. Goes with ``fork_from_app_id``.
+- ``fork_from_sequence_id`` -- Optional. The sequence_id to fork from. Goes with ``fork_from_app_id``.
+
 
-Note that you cannot use this in conjunction with :py:meth:`with_state <burr.core.application.ApplicationBuilder.with_state>`
+Note (1): that you cannot use this in conjunction with :py:meth:`with_state <burr.core.application.ApplicationBuilder.with_state>`
 or :py:meth:`with_entrypoint <burr.core.application.ApplicationBuilder.with_entrypoint>` -- these are mutually exclusive.
-Either you load from state or you start from scratch.
+
+Note (2): The loader will not error if no state is found, it will use the default state in this case.
+
+Forking State
+_____________
+Here's an explicit section on forking state. When loading you can also fork state from a previous application. This is useful if you want to start from a previous application's state,
+but don't want to add to it. The ``fork_from_app_id`` and ``fork_from_partition_key`` are used to identify the application to fork from, while
+``fork_from_sequence_id`` is used to identify the sequence_id to use. This is useful if you want to fork from a specific point in the application,
+rather than the latest state. This is especially useful for debugging, or building an application that enables you
+to rewind state and make different choices.
+
+When you use the ``fork_from_app_id`` to load state, the values passed to :py:meth:`with_identifiers() <burr.core.application.ApplicationBuilder.with_identifiers>`
+will then dictate where the new application state is ultimately stored.
+
 
 Writing state
 _____________