Spaces:
Runtime error
Runtime error
| """ | |
| State persistence utilities for flow manager. | |
| Provides saving flow state to local files or S3 storage. | |
| """ | |
| import json | |
| import os | |
| from datetime import datetime, timezone | |
| from typing import Any, Optional | |
| from loguru import logger | |
| from utils.s3_client import get_s3_client | |
| async def save_flow_state( | |
| flow_manager: Any, | |
| recording_dir: Optional[str] = None, | |
| s3_key: Optional[str] = None, | |
| use_local: bool = True, | |
| ) -> bool: | |
| """ | |
| Save flow manager state to either local file or S3. | |
| Args: | |
| flow_manager: The flow manager instance with state to save | |
| recording_dir: Directory for local recording (required if use_local=True) | |
| s3_key: S3 key for cloud storage (required if use_local=False) | |
| use_local: Whether to save locally or to S3 | |
| Returns: | |
| True if save succeeded, False otherwise | |
| """ | |
| if not flow_manager or not flow_manager.state: | |
| logger.debug("No flow manager state to save") | |
| return False | |
| logger.info("Saving flow manager state on cleanup") | |
| try: | |
| state_data = { | |
| "state": flow_manager.state, | |
| "current_node": flow_manager.current_node, | |
| "timestamp": datetime.now(timezone.utc).isoformat(), | |
| } | |
| state_json = json.dumps(state_data, indent=2) | |
| if use_local: | |
| return _save_to_local(state_json, recording_dir) | |
| else: | |
| return _save_to_s3(state_json, s3_key) | |
| except Exception as e: | |
| logger.error(f"Failed to save flow state: {e}") | |
| return False | |
| def _save_to_local(state_json: str, recording_dir: Optional[str]) -> bool: | |
| """Save state to local file.""" | |
| if not recording_dir: | |
| logger.error("Recording directory not provided for local save") | |
| return False | |
| state_file_path = os.path.join(recording_dir, "flow_state.json") | |
| try: | |
| os.makedirs(recording_dir, exist_ok=True) | |
| with open(state_file_path, "w") as f: | |
| f.write(state_json) | |
| logger.info(f"Flow state saved to {state_file_path}") | |
| return True | |
| except Exception as e: | |
| logger.error(f"Failed to save state to local file: {e}") | |
| return False | |
| def _save_to_s3(state_json: str, s3_key: Optional[str]) -> bool: | |
| """Save state to S3 using centralized S3Client.""" | |
| if not s3_key: | |
| logger.error("S3 key not provided for cloud save") | |
| return False | |
| try: | |
| s3_client = get_s3_client() | |
| return s3_client.upload_state(s3_key, state_json) | |
| except Exception as e: | |
| logger.error(f"Failed to save state to S3: {e}") | |
| return False | |