Spaces:
Runtime error
Runtime error
File size: 2,625 Bytes
24cdea4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """
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
|