"""State management for workflow execution.
This module provides tools for managing immutable state throughout workflow execution,
making it easier to handle state transitions in a predictable manner.
"""
import logging
from copy import deepcopy
from typing import Any, Generic, TypeVar, cast
from pydantic import BaseModel
logger = logging.getLogger(__name__)
# Type variable for the state model
StateT = TypeVar("StateT", bound=BaseModel)
[docs]
class StateManager:
"""Manages immutable state operations for workflow execution.
This class provides utilities for updating state objects immutably,
focusing on Pydantic models to ensure type safety and validation.
"""
[docs]
@staticmethod
def update_in(state_obj: BaseModel, path: list[str], value: Any) -> BaseModel:
"""Update a nested property in the state and return a new state object.
Args:
state_obj: The Pydantic model state object
path: List of attribute names forming a path to the property to update
value: The new value to set
Returns:
A new state object with the update applied
Raises:
TypeError: If state_obj is not a Pydantic BaseModel
KeyError: If the path is invalid
"""
if not isinstance(state_obj, BaseModel):
raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
# Create deep copy
new_state = state_obj.model_copy(deep=True)
# For simple top-level updates
if len(path) == 1:
setattr(new_state, path[0], value)
return new_state
# For nested updates
current = new_state
for i, key in enumerate(path[:-1]):
if not hasattr(current, key):
raise KeyError(f"Invalid path: {'.'.join(path[: i + 1])}")
# Get the next level object and ensure we're working with a copy
next_obj = getattr(current, key)
if isinstance(next_obj, BaseModel):
next_obj = next_obj.model_copy(deep=True)
setattr(current, key, next_obj)
elif isinstance(next_obj, dict) or isinstance(next_obj, list):
next_obj = deepcopy(next_obj)
setattr(current, key, next_obj)
current = next_obj
# Set the final value
if hasattr(current, path[-1]):
setattr(current, path[-1], value)
else:
raise KeyError(f"Invalid path: {'.'.join(path)}")
return new_state
[docs]
@staticmethod
def batch_update(
state_obj: BaseModel, updates: list[tuple[list[str], Any]]
) -> BaseModel:
"""Apply multiple updates to the state atomically.
Args:
state_obj: The Pydantic model state object
updates: List of (path, value) tuples with updates to apply
Returns:
A new state object with all updates applied
Raises:
TypeError: If state_obj is not a Pydantic BaseModel
KeyError: If any path is invalid
"""
if not isinstance(state_obj, BaseModel):
raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
# Create deep copy
new_state = state_obj.model_copy(deep=True)
# Apply each update
for path, value in updates:
new_state = StateManager.update_in(new_state, path, value)
return new_state
[docs]
@staticmethod
def get_in(state_obj: BaseModel, path: list[str]) -> Any:
"""Get the value at a nested path.
Args:
state_obj: The Pydantic model state object
path: List of attribute names forming a path to the property to retrieve
Returns:
The value at the specified path
Raises:
TypeError: If state_obj is not a Pydantic BaseModel
KeyError: If the path is invalid
"""
if not isinstance(state_obj, BaseModel):
raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
# For simple top-level properties
if len(path) == 1:
if not hasattr(state_obj, path[0]):
raise KeyError(f"Invalid path: {path[0]}")
return getattr(state_obj, path[0])
# For nested properties
current = state_obj
for i, key in enumerate(path):
if not hasattr(current, key):
raise KeyError(f"Invalid path: {'.'.join(path[: i + 1])}")
current = getattr(current, key)
return current
[docs]
@staticmethod
def merge(state_obj: BaseModel, **updates) -> BaseModel:
"""Merge flat updates into state and return a new state.
Args:
state_obj: The Pydantic model state object
**updates: Attribute updates to apply to the top level
Returns:
A new state object with the updates applied
Raises:
TypeError: If state_obj is not a Pydantic BaseModel
"""
if not isinstance(state_obj, BaseModel):
raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
return state_obj.model_copy(update=updates)
[docs]
class WorkflowStateWrapper(Generic[StateT]):
"""Wraps a state object with convenient update methods for use in workflows.
This wrapper provides a clean interface for immutable state updates
within workflow nodes, simplifying state management.
"""
[docs]
def __init__(self, state: StateT):
"""Initialize the state wrapper.
Args:
state: The Pydantic model state object to wrap
"""
self._state = state
[docs]
def update_in(self, path: list[str], value: Any) -> "WorkflowStateWrapper[StateT]":
"""Update state at path and return new wrapper.
Args:
path: List of attribute names forming a path to the property to update
value: The new value to set
Returns:
A new state wrapper with the update applied
"""
new_state = cast(StateT, StateManager.update_in(self._state, path, value))
return WorkflowStateWrapper(new_state)
[docs]
def batch_update(
self, updates: list[tuple[list[str], Any]]
) -> "WorkflowStateWrapper[StateT]":
"""Apply multiple updates to the state atomically.
Args:
updates: List of (path, value) tuples with updates to apply
Returns:
A new state wrapper with all updates applied
"""
new_state = cast(StateT, StateManager.batch_update(self._state, updates))
return WorkflowStateWrapper(new_state)
[docs]
def get_in(self, path: list[str]) -> Any:
"""Get the value at a nested path.
Args:
path: List of attribute names forming a path to the property to retrieve
Returns:
The value at the specified path
"""
return StateManager.get_in(self._state, path)
[docs]
def merge(self, **updates) -> "WorkflowStateWrapper[StateT]":
"""Merge flat updates into state and return a new wrapper.
Args:
**updates: Attribute updates to apply to the top level
Returns:
A new state wrapper with the updates applied
"""
new_state = cast(StateT, StateManager.merge(self._state, **updates))
return WorkflowStateWrapper(new_state)
[docs]
def get_state(self) -> StateT:
"""Get the wrapped state object.
Returns:
The current state object
"""
return self._state
[docs]
def __repr__(self) -> str:
"""Get string representation."""
return f"WorkflowStateWrapper({self._state})"