diff --git a/mephisto/abstractions/blueprints/parlai_chat/parlai_chat_agent_state.py b/mephisto/abstractions/blueprints/parlai_chat/parlai_chat_agent_state.py index 9a0004b2e..fac2d1a4e 100644 --- a/mephisto/abstractions/blueprints/parlai_chat/parlai_chat_agent_state.py +++ b/mephisto/abstractions/blueprints/parlai_chat/parlai_chat_agent_state.py @@ -57,7 +57,15 @@ def _load_data(self) -> None: self.messages = state["outputs"]["messages"] self.init_data = state["inputs"] self.final_submission = state["outputs"].get("final_submission") - self.metadata = _AgentStateMetadata(**state["metadata"]) + if "metadata" in state: + self.metadata = _AgentStateMetadata(**state["metadata"]) + elif "times" in state: + self.metadata = _AgentStateMetadata( + start_time=state["times"]["start_time"], + end_time=state["times"]["end_time"], + ) + else: + self.metadata = _AgentStateMetadata() def get_data(self) -> Dict[str, Any]: """Return dict with the messages of this agent"""