Save agent state (#7372)

This commit is contained in:
Engel Nyst 2025-03-20 05:16:49 +01:00 committed by GitHub
parent f5aeb47a72
commit 34920ea04e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -111,12 +111,13 @@ class State:
get_conversation_agent_state_filename(sid, user_id), encoded
)
# see if state is in old directory. If yes, delete it.
filename = get_conversation_agent_state_filename(sid)
try:
file_store.delete(filename)
except Exception:
pass
# see if state is in the old directory on saas/remote use cases and delete it.
if user_id:
filename = get_conversation_agent_state_filename(sid)
try:
file_store.delete(filename)
except Exception:
pass
except Exception as e:
logger.error(f'Failed to save state to session: {e}')
raise e
@ -125,6 +126,11 @@ class State:
def restore_from_session(
sid: str, file_store: FileStore, user_id: str | None = None
) -> 'State':
"""
Restores the state from the previously saved session.
"""
state: State
try:
encoded = file_store.read(
get_conversation_agent_state_filename(sid, user_id)
@ -132,12 +138,17 @@ class State:
pickled = base64.b64decode(encoded)
state = pickle.loads(pickled)
except FileNotFoundError:
# if user_id is provided, we are in a saas/remote use case
# and we need to check if the state is in the old directory.
if user_id:
# see if state is in old directory. If yes, load it.
filename = get_conversation_agent_state_filename(sid)
encoded = file_store.read(filename)
pickled = base64.b64decode(encoded)
state = pickle.loads(pickled)
else:
raise FileNotFoundError(
f'Could not restore state from session file for sid: {sid}'
)
except Exception as e:
logger.debug(f'Could not restore state from session: {e}')
raise e