style: change of formatting

Author: jmatejcz

src/rai_bench/rai_bench/benchmark_model.py

@@ -28,7 +28,9 @@
     SpawnedEntity,
 )
 
-SimulationBridgeT = TypeVar("SimulationBridgeT", bound=SimulationBridge)
+SimulationBridgeT = TypeVar(
+    "SimulationBridgeT", bound=SimulationBridge[SimulationConfig]
+)
 loggers_type = Union[RcutilsLogger, logging.Logger]
 
 
@@ -70,7 +72,9 @@ def validate_config(self, simulation_config: SimulationConfig) -> bool:
         pass
 
     @abstractmethod
-    def calculate_result(self, simulation_bridge: SimulationBridge) -> float:
+    def calculate_result(
+        self, simulation_bridge: SimulationBridge[SimulationConfig]
+    ) -> float:
         """
         Calculate result of the task
         """
@@ -95,7 +99,7 @@ def is_adjacent(self, pos1: PoseModel, pos2: PoseModel, threshold_distance: floa
         Check if positions are adjacent to each other, the threshold_distance is a distance
         in simulation, refering to how close they have to be to classify them as adjacent
         """
-        self.logger.debug(
+        self.logger.debug(  # type: ignore
             f"Euclidean distance: {self.euclidean_distance(pos1, pos2)}, pos1: {pos1}, pos2: {pos2}"
         )
         return self.euclidean_distance(pos1, pos2) < threshold_distance
@@ -164,7 +168,7 @@ def __init__(
     def create_scenarios(
         cls, tasks: List[Task], simulation_configs: List[SimulationConfig]
     ) -> list[Any]:
-        scenarios = []
+        scenarios: List[Scenario] = []
         for task in tasks:
             for sim_conf in simulation_configs:
                 try:
@@ -175,22 +179,22 @@ def create_scenarios(
                     )
         return scenarios
 
-    def run_next(self, agent):
+    def run_next(self, agent) -> None:
         """
         Runs the next scenario
         """
         try:
             i, scenario = next(self.scenarios)  # Get the next scenario
 
             self.simulation_bridge.setup_scene(scenario.simulation_config)
-            self._logger.info(
+            self._logger.info(  # type: ignore
                 "======================================================================================"
             )
-            self._logger.info(
+            self._logger.info(  # type: ignore
                 f"RUNNING SCENARIO NUMBER {i + 1}, TASK: {scenario.task.get_prompt()}"
             )
             initial_result = scenario.task.calculate_result(self.simulation_bridge)
-            self._logger.info(f"RESULT OF THE INITIAL SETUP: {initial_result}")
+            self._logger.info(f"RESULT OF THE INITIAL SETUP: {initial_result}")  # type: ignore
             tool_calls_num = 0
 
             ts = time.perf_counter()
@@ -209,7 +213,7 @@ def run_next(self, agent):
                                 last_msg = msg.content[0].get("text", "")
                     else:
                         last_msg = msg.content
-                        self._logger.debug(f"{graph_node_name}: {last_msg}")
+                        self._logger.debug(f"{graph_node_name}: {last_msg}")  # type: ignore
 
                 else:
                     raise ValueError(f"Unexpected type of message: {type(msg)}")
@@ -218,13 +222,13 @@ def run_next(self, agent):
                     # TODO (jm) figure out more robust way of counting tool calls
                     tool_calls_num += len(msg.tool_calls)
 
-                self._logger.info(f"AI Message: {msg}")
+                self._logger.info(f"AI Message: {msg}")  # type: ignore
 
             te = time.perf_counter()
 
             result = scenario.task.calculate_result(self.simulation_bridge)
             total_time = te - ts
-            self._logger.info(
+            self._logger.info(  # type: ignore
                 f"TASK SCORE: {result}, TOTAL TIME: {total_time:.3f}, NUM_OF_TOOL_CALLS: {tool_calls_num}"
             )
 
@@ -241,5 +245,5 @@ def run_next(self, agent):
         except StopIteration:
             print("No more scenarios left to run.")
 
-    def get_results(self) -> list[dict]:
+    def get_results(self) -> List[Dict[str, Any]]:
         return self.results

src/rai_bench/rai_bench/main.py

@@ -16,26 +16,25 @@
 import logging
 import time
 from pathlib import Path
+from typing import List
 
 import rclpy
-import rclpy.qos
+from langchain.tools import BaseTool
 from rai_open_set_vision.tools import GetGrabbingPointTool
 
 from rai.agents.conversational_agent import create_conversational_agent
 from rai.communication.ros2.connectors import ROS2ARIConnector
 from rai.tools.ros.manipulation import GetObjectPositionsTool, MoveToPointTool
 from rai.tools.ros2.topics import GetROS2ImageTool, GetROS2TopicsNamesAndTypesTool
 from rai.utils.model_initialization import get_llm_model
-from rai_bench.benchmark_model import (
-    Benchmark,
-)
+from rai_bench.benchmark_model import Benchmark, Task
 from rai_bench.o3de_test_bench.tasks import GrabCarrotTask, PlaceCubesTask
 from rai_sim.o3de.o3de_bridge import (
     O3DEngineArmManipulationBridge,
     O3DExROS2SimulationConfig,
     PoseModel,
 )
-from rai_sim.simulation_bridge import Rotation, Translation
+from rai_sim.simulation_bridge import Rotation, SimulationConfig, Translation
 
 if __name__ == "__main__":
     rclpy.init()
@@ -55,7 +54,7 @@
     Before starting the task, make sure to grab the camera image to understand the environment.
     """
     # define tools
-    tools = [
+    tools: List[BaseTool] = [
         GetObjectPositionsTool(
             connector=connector,
             target_frame="panda_link0",
@@ -80,7 +79,7 @@
     file_handler.setFormatter(formatter)
 
     bench_logger = logging.getLogger("Benchmark logger")
-    bench_logger.setLevel(logging.DEBUG)
+    bench_logger.setLevel(logging.INFO)
     bench_logger.addHandler(file_handler)
 
     agent_logger = logging.getLogger("Agent logger")
@@ -134,11 +133,11 @@
         configs_dir + "scene3.yaml",
         configs_dir + "scene4.yaml",
     ]
-    simulations_configs = [
+    simulations_configs: List[SimulationConfig] = [
         O3DExROS2SimulationConfig.load_config(Path(path), Path(connector_path))
         for path in scene_paths
     ]
-    tasks = [
+    tasks: List[Task] = [
         GrabCarrotTask(logger=bench_logger),
         PlaceCubesTask(logger=bench_logger),
     ]

src/rai_bench/rai_bench/o3de_test_bench/tasks/grab_carrot_task.py

@@ -33,7 +33,9 @@ def validate_config(self, simulation_config: SimulationConfig) -> bool:
 
         return False
 
-    def calculate_result(self, simulation_bridge: SimulationBridge) -> float:
+    def calculate_result(
+        self, simulation_bridge: SimulationBridge[SimulationConfig]
+    ) -> float:
         # TODO (jm) extract common logic to some parent manipulation task?
         initially_misplaced_now_correct = 0  # when the object which was in the incorrect place at the start, is in a correct place at the end
         initially_misplaced_still_incorrect = 0  # when the object which was in the incorrect place at the start, is in a incorrect place at the end
@@ -55,8 +57,8 @@ def calculate_result(self, simulation_bridge: SimulationBridge) -> float:
             )
 
         else:
-            self.logger.debug(f"initial positions: {initial_carrots}")
-            self.logger.debug(f"current positions: {final_carrots}")
+            self.logger.debug(f"initial positions: {initial_carrots}")  # type: ignore
+            self.logger.debug(f"current positions: {final_carrots}")  # type: ignore
             for ini_carrot in initial_carrots:
                 for final_carrot in final_carrots:
                     if ini_carrot.name == final_carrot.name:
@@ -91,7 +93,7 @@ def calculate_result(self, simulation_bridge: SimulationBridge) -> float:
                         f"Entity with name: {ini_carrot.name} which was present in initial scene, not found in final scene."
                     )
 
-            self.logger.info(
+            self.logger.info(  # type: ignore
                 f"initially_misplaced_now_correct: {initially_misplaced_now_correct}, initially_misplaced_still_incorrect: {initially_misplaced_still_incorrect}, initially_correct_still_correct: {initially_correct_still_correct}, initially_correct_now_incorrect: {initially_correct_now_incorrect}"
             )
             return (