Backward compatibility with bgym stepInfo

Author: amanjaiswal73892

src/agentlab/analyze/agent_xray.py

@@ -28,6 +28,7 @@
 from agentlab.llm.llm_utils import BaseMessage as AgentLabBaseMessage
 from agentlab.llm.llm_utils import Discussion
 from agentlab.llm.response_api import MessageBuilder, ToolCalls
+from browsergym.experiments.loop import StepInfo as BGymStepInfo
 
 select_dir_instructions = "Select Experiment Directory"
 AGENT_NAME_KEY = "agent.agent_name"
@@ -1319,7 +1320,7 @@ def plot_profiling(ax, step_info_list: list[StepInfo], summary_info: dict, progr
     all_times = []
     step_times = []
     for i, step_info in progress_fn(list(enumerate(step_info_list)), desc="Building plot."):
-        assert isinstance(step_info, StepInfo), f"Expected StepInfo, got {type(step_info)}"
+        assert isinstance(step_info, (StepInfo, BGymStepInfo)), f"Expected StepInfo or BGymStepInfo, got {type(step_info)}"
         step = step_info.step
 
         prof = deepcopy(step_info.profiling)