fix: collecting messages from agent

src/rai_bench/rai_bench/manipulation_o3de/benchmark.py

@@ -184,33 +184,36 @@ def run_next(self, agent: CompiledStateGraph) -> None:
             tool_calls_num = 0
 
             ts = time.perf_counter()
+            prev_count: int = 0
             for state in agent.stream(
                 {"messages": [HumanMessage(content=scenario.task.get_prompt())]},
                 {
                     "recursion_limit": 100
                 },  # NOTE (jmatejcz) what should be recursion limit?
             ):
-                graph_node_name = list(state.keys())[0]
-                msg = state[graph_node_name]["messages"][-1]
-
-                if isinstance(msg, HumanMultimodalMessage):
-                    last_msg = msg.text
-                elif isinstance(msg, BaseMessage):
-                    if isinstance(msg.content, list):
-                        if len(msg.content) == 1:
-                            if type(msg.content[0]) is dict:
-                                last_msg = msg.content[0].get("text", "")
-                    else:
-                        last_msg = msg.content
-                        self._logger.debug(f"{graph_node_name}: {last_msg}")  # type: ignore
+                node = next(iter(state))
+                new_messages = state[node]["messages"][prev_count:]
+                prev_count = len(state[node]["messages"])
+
+                for msg in new_messages:
+                    if isinstance(msg, HumanMultimodalMessage):
+                        last_msg = msg.text
+                    elif isinstance(msg, BaseMessage):
+                        if isinstance(msg.content, list):
+                            if len(msg.content) == 1:
+                                if type(msg.content[0]) is dict:
+                                    last_msg = msg.content[0].get("text", "")
+                        else:
+                            last_msg = msg.content
+                            self._logger.debug(f"{node}: {last_msg}")  # type: ignore
 
-                else:
-                    raise ValueError(f"Unexpected type of message: {type(msg)}")
+                    else:
+                        raise ValueError(f"Unexpected type of message: {type(msg)}")
 
-                if isinstance(msg, AIMessage):
-                    tool_calls_num += len(msg.tool_calls)
+                    if isinstance(msg, AIMessage):
+                        tool_calls_num += len(msg.tool_calls)
 
-                self._logger.info(f"AI Message: {msg}")  # type: ignore
+                    self._logger.info(f"AI Message: {msg}")  # type: ignore
 
             te = time.perf_counter()
             try:

src/rai_bench/rai_bench/tool_calling_agent/benchmark.py

@@ -181,9 +181,10 @@ def run_next(self, agent: CompiledStateGraph, model_name: str) -> None:
 
         ts = time.perf_counter()
         messages: List[BaseMessage] = []
+        prev_count: int = 0
         try:
             if isinstance(task, SpatialReasoningAgentTask):
-                for event in agent.stream(
+                for state in agent.stream(
                     {
                         "messages": [
                             HumanMultimodalMessage(
@@ -193,20 +194,24 @@ def run_next(self, agent: CompiledStateGraph, model_name: str) -> None:
                     },
                     config=config,
                 ):
-                    flattened = {k: v for d in event.values() for k, v in d.items()}
-                    # cos = event.values()
-                    messages.extend(flattened["messages"])
+                    node = next(iter(state))
+                    all_messages = state[node]["messages"]
+                    for new_msg in all_messages[prev_count:]:
+                        messages.append(new_msg)
+                    prev_count = len(messages)
             else:
-                for event in agent.stream(
+                for state in agent.stream(
                     {"messages": [HumanMultimodalMessage(content=task.get_prompt())]},
                     config=config,
                 ):
-                    flattened = {k: v for d in event.values() for k, v in d.items()}
-                    messages.extend(flattened["messages"])
+                    node = next(iter(state))
+                    all_messages = state[node]["messages"]
+                    for new_msg in all_messages[prev_count:]:
+                        messages.append(new_msg)
+                    prev_count = len(messages)
 
         except GraphRecursionError as e:
             self.logger.error(msg=f"Reached recursion limit {e}")
-            # task.fail_rest_of_validators()
 
         self.logger.debug(messages)
         toll_calls = task.get_tool_calls_from_messages(messages=messages)
@@ -252,11 +257,6 @@ def run_next(self, agent: CompiledStateGraph, model_name: str) -> None:
         if completed_tasks == self.num_tasks:
             self._compute_and_save_summary()
 
-        # except StopIteration:
-        #     if self.task_results:
-        #         self._compute_and_save_summary()
-        #     print("No more scenarios left to run.")
-
     def _compute_and_save_summary(self):
         self.logger.info("Computing and saving average results...")
         for model_name, results in self.model_results.items():

src/rai_extensions/rai_open_set_vision/scripts/run_vision_agents.py

@@ -14,7 +14,7 @@
 
 
 import rclpy
-from rai.utils import wait_for_shutdown
+from rai.agents import wait_for_shutdown
 from rai_open_set_vision.agents import GroundedSamAgent, GroundingDinoAgent