feat: add StateBaseAgent
#529
Conversation
boczekbartek
force-pushed
the
bb/feat/state_based_agent
branch
from
66e2383 to
70c68a3
Compare
boczekbartek
force-pushed
the
bb/feat/state_based_agent
branch
from
b7ad8a6 to
fc10dcb
Compare
|
@coderabbitai full review |
WalkthroughThis update introduces a modular framework for state-based agents with ROS2 integration, focusing on multimodal message aggregation and stateful language model interactions. It adds new aggregator abstractions and ROS2-specific aggregators for logs and images, and implements a generic Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant ROS2StateBasedAgent
participant Aggregators
participant ROS2Connector
participant LLM
participant Tools
User->>ROS2StateBasedAgent: Start agent
ROS2StateBasedAgent->>ROS2Connector: Setup connectors (subscribe/publish)
loop Periodic Aggregation
ROS2Connector->>Aggregators: Receive and buffer messages
Aggregators->>ROS2StateBasedAgent: Aggregate state (logs, images, etc.)
end
ROS2StateBasedAgent->>LLM: Provide aggregated state (via retriever)
LLM->>Tools: (Conditional) Call tools if needed
Tools-->>LLM: Tool results
LLM->>ROS2StateBasedAgent: Generate response
ROS2StateBasedAgent->>ROS2Connector: Publish response (e.g., to HRI topic)
User-->>ROS2StateBasedAgent: Shutdown signal
ROS2StateBasedAgent->>ROS2Connector: Clean shutdown
sequenceDiagram
participant Agent
participant Aggregators
participant LLM
loop Aggregation Interval
Agent->>Aggregators: Trigger aggregation
Aggregators-->>Agent: Aggregated messages (logs, images, etc.)
end
Agent->>LLM: Run stateful runnable with aggregated state
LLM->>Agent: Output (response, tool calls, etc.)
📜 Recent review detailsConfiguration used: .coderabbit.yaml 📒 Files selected for processing (16)
💤 Files with no reviewable changes (2)
🧰 Additional context used🧬 Code Graph Analysis (6)src/rai_core/rai/agents/ros2/state_based_agent.py (2)
src/rai_core/rai/aggregators/__init__.py (1)
src/rai_core/rai/agents/ros2/__init__.py (1)
src/rai_core/rai/aggregators/base.py (3)
src/rai_core/rai/communication/ros2/api/conversion.py (1)
src/rai_core/rai/agents/langchain/agent.py (3)
🔇 Additional comments (13)
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🔭 Outside diff range comments (1)
src/rai_core/rai/communication/ros2/api/conversion.py (1)
117-134: 🛠️ Refactor suggestionConsider consolidating image conversion functions.
The new function provides a cleaner approach for encoding ROS2 images to base64, supporting both
ImageandCompressedImagetypes. It's a good implementation, but as noted in the TODO comment, it should be merged with the existingconvert_ros_img_to_base64function to avoid code duplication.Consider implementing this consolidation by either:
- Expanding the existing function to handle both types
- Creating a unified function that calls appropriate conversion methods based on message type
-# TODO(boczekbartek): merge with above -def encode_ros2_img_to_base64( - img_message: sensor_msgs.msg.Image | sensor_msgs.msg.CompressedImage, -) -> str: - msg_type = type(img_message) - if msg_type == sensor_msgs.msg.Image: - image = CvBridge().imgmsg_to_cv2( # type: ignore - img_message, desired_encoding="rgb8" - ) - elif msg_type == sensor_msgs.msg.CompressedImage: - image = CvBridge().compressed_imgmsg_to_cv2( # type: ignore - img_message, desired_encoding="rgb8" - ) - else: - raise ValueError(f"Unsupported message type: {msg_type}") - image = cast(cv2.Mat, image) - return preprocess_image(image) +def convert_ros_img_to_base64( + msg: sensor_msgs.msg.Image | sensor_msgs.msg.CompressedImage, + desired_encoding: str = "rgb8" +) -> str: + """Convert ROS2 image message to base64 encoded string. + + Args: + msg: ROS2 image message + desired_encoding: Desired encoding for the image + + Returns: + Base64 encoded string + + Raises: + ValueError: If the message type is not supported + """ + bridge = CvBridge() + msg_type = type(msg) + + if msg_type == sensor_msgs.msg.Image: + cv_image = cast(cv2.Mat, bridge.imgmsg_to_cv2(msg, desired_encoding=desired_encoding)) + elif msg_type == sensor_msgs.msg.CompressedImage: + cv_image = cast(cv2.Mat, bridge.compressed_imgmsg_to_cv2(msg, desired_encoding=desired_encoding)) + else: + raise ValueError(f"Unsupported message type: {msg_type}") + + return preprocess_image(cv_image)
🧹 Nitpick comments (5)
examples/agents/state_based.py (1)
12-12: Typo in License HeaderThere's a small spelling error in line 12: "goveself.rning" should be "governing".
-# See the License for the specific language goveself.rning permissions and +# See the License for the specific language governing permissions andsrc/rai_core/rai/agents/react_agent.py (1)
39-49: Clarify Handling of 'runnable' vs 'tools' and 'system_prompt'When both
runnableandtools/system_promptare provided, the latter are effectively superseded by the customrunnable. Consider documenting or warning if both are passed, to avoid confusion.src/rai_core/rai/agents/langchain/runnables.py (2)
126-127: Add Docstring to ExplainStateMessageA short docstring describing the purpose and usage of
StateMessagecan aid readability and maintainability.+class StateMessage(HumanMultimodalMessage): + """ + A multimodal message containing retrieved state information + (images, audio, textual content) for the agent to incorporate. + """
162-164: Use Named Function Instead of LambdaStatic analysis suggests rewriting the lambda expression as a named function to enhance clarity.
- state_retriever = lambda: {} + def default_state_retriever(): + return {} + state_retriever = default_state_retriever🧰 Tools
🪛 Ruff (0.8.2)
163-163: Do not assign a
lambdaexpression, use adefRewrite
state_retrieveras adef(E731)
src/rai_core/rai/state_based/ros2/aggregators.py (1)
98-99: Optimize the image encoding step.Currently, you encode all images in the buffer even though only the last image is used in the aggregator. This can be inefficient for large buffers or high-frequency image streams.
-def __call__(self, msgs: List[Image | CompressedImage]) -> str | None: - if len(msgs) == 0: - return None - b64_images: List[str] = [encode_ros2_img_to_base64(msg) for msg in msgs] - return encode_ros2_img_to_base64(msgs[-1]) +def __call__(self, msgs: List[Image | CompressedImage]) -> str | None: + if not msgs: + return None + return encode_ros2_img_to_base64(msgs[-1])
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
examples/agents/state_based.py(1 hunks)src/rai_core/rai/agents/__init__.py(1 hunks)src/rai_core/rai/agents/langchain/runnables.py(2 hunks)src/rai_core/rai/agents/react_agent.py(2 hunks)src/rai_core/rai/agents/state_based.py(0 hunks)src/rai_core/rai/agents/state_based_agent.py(1 hunks)src/rai_core/rai/communication/ros2/api/conversion.py(2 hunks)src/rai_core/rai/initialization/model_initialization.py(2 hunks)src/rai_core/rai/state_based/ros2/aggregators.py(1 hunks)src/rai_core/rai/tools/ros2/generic/topics.py(2 hunks)
💤 Files with no reviewable changes (1)
- src/rai_core/rai/agents/state_based.py
🧰 Additional context used
🧬 Code Graph Analysis (5)
src/rai_core/rai/tools/ros2/generic/topics.py (2)
src/rai_core/rai/communication/ros2/api/conversion.py (2)
encode_ros2_img_to_base64(118-133)ros2_message_to_dict(32-47)src/rai_core/rai/messages/artifacts.py (1)
MultimodalArtifact(20-22)
src/rai_core/rai/agents/react_agent.py (1)
src/rai_core/rai/agents/langchain/runnables.py (2)
ReActAgentState(38-47)create_react_runnable(78-123)
src/rai_core/rai/communication/ros2/api/conversion.py (1)
src/rai_core/rai/messages/conversion.py (1)
preprocess_image(24-48)
src/rai_core/rai/agents/__init__.py (1)
src/rai_core/rai/agents/state_based_agent.py (2)
StateBasedAgent(45-166)StateBasedConfig(36-42)
src/rai_core/rai/agents/langchain/runnables.py (3)
src/rai_core/rai/agents/tool_runner.py (1)
ToolRunner(33-146)src/rai_core/rai/initialization/model_initialization.py (1)
get_llm_model(125-158)src/rai_core/rai/messages/multimodal.py (1)
HumanMultimodalMessage(60-81)
🪛 Ruff (0.8.2)
src/rai_core/rai/agents/langchain/runnables.py
163-163: Do not assign a lambda expression, use a def
Rewrite state_retriever as a def
(E731)
src/rai_core/rai/state_based/ros2/aggregators.py
51-51: Do not perform function call TopicConfig in argument defaults; instead, perform the call within the function, or read the default from a module-level singleton variable
(B008)
🔇 Additional comments (7)
src/rai_core/rai/communication/ros2/api/conversion.py (1)
29-29: Good addition of import for preprocess_image.The import statement correctly brings in the preprocess_image function that will be used by the new encode_ros2_img_to_base64 function.
src/rai_core/rai/agents/__init__.py (2)
18-18: Updated import for the new StateBasedAgent architecture.The import statement has been properly updated to reflect the refactoring from the factory function pattern to a class-based approach.
24-25: Properly updated all list with new exported symbols.The all list has been correctly updated to include the new StateBasedAgent and StateBasedConfig classes, which aligns with the architectural changes in this PR.
src/rai_core/rai/initialization/model_initialization.py (2)
22-22: Good addition of BaseChatModel import.Importing BaseChatModel from langchain.chat_models.base is a necessary addition to support the return type annotation being added.
130-130: Enhanced type safety with return type annotation.Adding the explicit return type annotation -> BaseChatModel to the get_llm_model function improves type safety and code readability, making it clearer what type of object the function returns.
src/rai_core/rai/tools/ros2/generic/topics.py (2)
25-29: Well-structured imports with specific functions.The imports have been properly updated to include only the necessary functions from the conversion module, improving code clarity by making dependencies explicit.
136-137:✅ Verification successful
Simplified image processing using centralized utility function.
The implementation now leverages the new
encode_ros2_img_to_base64utility function, which reduces code duplication and centralizes image processing logic. This change makes the code more maintainable and consistent.
🏁 Script executed:
#!/bin/bash # Verify that no other code in the repository still uses direct CvBridge conversion # for images that could benefit from the new utility function echo "Searching for potential direct CvBridge usage for image conversion that could use encode_ros2_img_to_base64 instead:" rg -A 3 -B 3 "CvBridge\(\).imgmsg_to_cv2|CvBridge\(\).compressed_imgmsg_to_cv2" --type py --glob '!**/conversion.py'Length of output: 361
Centralized image processing utility verified
The changes introduced in
src/rai_core/rai/tools/ros2/generic/topics.py(lines 136-137) now effectively leverage theencode_ros2_img_to_base64utility function. The provided grep command confirmed there are no residual direct usages of CvBridge conversion methods in the repository, ensuring consistency and improved maintainability.
There was a problem hiding this comment.
Consider renaming state_based catalog to aggregators.
Consider moving the BaseStateAggregator(ABC) to base.py and keeping the implementations in aggregators/ros2 catalog.
There was a problem hiding this comment.
good idea, done in cfd0ca0
| qos_profile = self._ari_connector._topic_api._resolve_qos_profile( | ||
| topic, config.auto_qos_matching, config.qos_profile, for_publisher=False | ||
| ) | ||
| print(qos_profile) |
There was a problem hiding this comment.
| print(qos_profile) |
There was a problem hiding this comment.
sorry for not accepting your suggestion, I noticed your comment after cleaning this up myself :D
boczekbartek
force-pushed
the
bb/feat/state_based_agent
branch
from
4c06d7c to
f31ffff
Compare
boczekbartek
force-pushed
the
bb/feat/state_based_agent
branch
from
a0665a3 to
6620b74
Compare
|
@coderabbitai full review |
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (10)
src/rai_core/rai/aggregators/base.py (1)
49-56: Consider adding return type annotations to utility methods.While the code is well-structured, for consistency with the rest of the class, consider adding return type annotations to the
clearandget_buffermethods.- def clear(self) -> None: + def clear(self) -> None: self._buffer.clear() - def get_buffer(self) -> List[T]: + def get_buffer(self) -> List[T]: return list(self._buffer) - def __str__(self) -> str: + def __str__(self) -> str: return f"{self.__class__.__name__}"src/rai_core/rai/agents/ros2/state_based_agent.py (1)
19-21: Consider adding constructor parameters for ROS2Connector customization.The implementation is simple and follows the inheritance pattern correctly. However, the
ROS2Connectorconstructor accepts parameters likenode_nameanddestroy_subscribersthat could be useful for customization.Consider extending the constructor to allow customization of the ROS2 connector:
-class ROS2StateBasedAgent(BaseStateBasedAgent): - def setup_connector(self): - return ROS2Connector() +class ROS2StateBasedAgent(BaseStateBasedAgent): + def __init__( + self, + connectors: dict[str, HRIConnector[HRIMessage]], + config: StateBasedConfig, + llm: Optional[BaseChatModel] = None, + tools: Optional[List[BaseTool]] = None, + state: Optional[ReActAgentState] = None, + system_prompt: Optional[str] = None, + node_name: str = None, + destroy_subscribers: bool = False, + ): + super().__init__( + connectors, config, llm, tools, state, system_prompt + ) + self.node_name = node_name + self.destroy_subscribers = destroy_subscribers + + def setup_connector(self): + return ROS2Connector( + node_name=self.node_name, + destroy_subscribers=self.destroy_subscribers + )This allows users to customize the ROS2 connector when creating the agent.
src/rai_core/rai/tools/ros2/generic/topics.py (1)
136-137: Simplified image encoding using the new utility function.The code is much cleaner now, using the centralized encoding function instead of inline conversion logic. However, consider adding error handling for potential failures during image encoding.
- b64_image = encode_ros2_img_to_base64(message.payload) - return "Image received successfully", MultimodalArtifact(images=[b64_image]) + try: + b64_image = encode_ros2_img_to_base64(message.payload) + return "Image received successfully", MultimodalArtifact(images=[b64_image]) + except ValueError as e: + # Handle unsupported message types + return f"Error encoding image: {str(e)}", MultimodalArtifact(images=[]) + except Exception as e: + # Handle other errors + return f"Unexpected error encoding image: {str(e)}", MultimodalArtifact(images=[])src/rai_core/rai/agents/react_agent.py (1)
39-49:self.llmmay stayNonewhen a pre‑built runnable is injected – confirm this will not break downstream usagesWhen the caller passes in a ready‑made
runnable, thellmargument can legitimately beNone.
Because we unconditionally store it (self.llm = llm), any later code that expects a concreteBaseChatModelmay findNoneand raise.If
self.llmis never referenced outside this constructor, consider removing the attribute entirely or documenting that it can beNone. Otherwise, guard all later usages.examples/agents/state_based.py (1)
11-13: Typo in licence header
goveself.rning→governingWhile purely cosmetic, licence headers are scanned automatically by some compliance tools; correcting the typo avoids false positives.
-# See the License for the specific language goveself.rning permissions and +# See the License for the specific language governing permissions andsrc/rai_core/rai/agents/base_state_based_agent.py (2)
80-88: UnusedReentrantCallbackGroup—remove or integrate to avoid dead code
self._callback_groupis created but never passed to the connector or used elsewhere.
Keeping unused members adds cognitive overhead and may mis‑lead future maintainers.- self._callback_group = ReentrantCallbackGroup()If you plan to leverage it for ROS2 subscriptions later, add a
TODOcomment explaining when it will become relevant. Otherwise, delete it.
154-165:_stop_event.clear()at the end ofstop()can restart threads unintentionally
clear()resets the flag that other threads (andrun()) rely on.
If a user callsstop()and later forgets to recreate the agent, a subsequentrun()will restart old threads on the same object, which is error‑prone.Remove the
clear()or add a comment justifying the reuse scenario.src/rai_core/rai/aggregators/ros2/aggregators.py (3)
34-50: Consider clearing the buffer after processingWithout
self.clear(), the same log messages are re‑processed in every interval.
If persistence is not desired, addself.clear()at the end ofget().
130-157: Diff aggregator: misleading statistics and missing original lengthThe final message uses
len(b64_images)twice, so it always reports n keyframes selected from n images.- return HumanMessage( - content=f"Result of the analysis of the {len(b64_images)} keyframes selected from {len(b64_images)} last images:\n{response}" - ) + return HumanMessage( + content=( + f"Result of the analysis of the {len(b64_images)} keyframes " + f"selected from {len(msgs)} last images:\n{response}" + ) + )
111-124:get_key_elementsdocstring & name mention 1st/middle/last but function returns in unspecified orderThe helper currently returns
[elements[0], elements[middle], elements[-1]], which is fine, but consider clarifying the order in the docstring or renaming toselect_keyframesfor intent clarity.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Lite
📒 Files selected for processing (15)
examples/agents/state_based.py(1 hunks)src/rai_core/rai/agents/__init__.py(1 hunks)src/rai_core/rai/agents/base_state_based_agent.py(1 hunks)src/rai_core/rai/agents/langchain/runnables.py(2 hunks)src/rai_core/rai/agents/react_agent.py(2 hunks)src/rai_core/rai/agents/ros2/__init__.py(1 hunks)src/rai_core/rai/agents/ros2/state_based_agent.py(1 hunks)src/rai_core/rai/agents/state_based.py(0 hunks)src/rai_core/rai/aggregators/__init__.py(1 hunks)src/rai_core/rai/aggregators/base.py(1 hunks)src/rai_core/rai/aggregators/ros2/__init__.py(1 hunks)src/rai_core/rai/aggregators/ros2/aggregators.py(1 hunks)src/rai_core/rai/communication/ros2/api/conversion.py(2 hunks)src/rai_core/rai/initialization/model_initialization.py(2 hunks)src/rai_core/rai/tools/ros2/generic/topics.py(2 hunks)
💤 Files with no reviewable changes (1)
- src/rai_core/rai/agents/state_based.py
🧰 Additional context used
🧬 Code Graph Analysis (8)
src/rai_core/rai/tools/ros2/generic/topics.py (2)
src/rai_core/rai/communication/ros2/api/conversion.py (2)
encode_ros2_img_to_base64(118-133)ros2_message_to_dict(32-47)src/rai_core/rai/messages/artifacts.py (1)
MultimodalArtifact(20-22)
src/rai_core/rai/aggregators/__init__.py (1)
src/rai_core/rai/aggregators/base.py (1)
BaseAggregator(26-56)
src/rai_core/rai/aggregators/ros2/__init__.py (1)
src/rai_core/rai/aggregators/ros2/aggregators.py (4)
ROS2GetLastImageAggregator(54-63)ROS2ImgVLMDescriptionAggregator(66-102)ROS2ImgVLMDiffAggregator(105-157)ROS2LogsAggregator(28-51)
src/rai_core/rai/agents/ros2/state_based_agent.py (2)
src/rai_core/rai/agents/base_state_based_agent.py (2)
BaseStateBasedAgent(47-165)setup_connector(93-94)src/rai_core/rai/communication/ros2/connectors/connector.py (1)
ROS2Connector(42-283)
src/rai_core/rai/agents/ros2/__init__.py (1)
src/rai_core/rai/agents/ros2/state_based_agent.py (1)
ROS2StateBasedAgent(19-21)
src/rai_core/rai/communication/ros2/api/conversion.py (1)
src/rai_core/rai/messages/conversion.py (1)
preprocess_image(24-48)
src/rai_core/rai/agents/react_agent.py (1)
src/rai_core/rai/agents/langchain/runnables.py (2)
ReActAgentState(39-48)create_react_runnable(79-124)
src/rai_core/rai/aggregators/base.py (1)
src/rai_core/rai/messages/multimodal.py (1)
HumanMultimodalMessage(60-81)
🔇 Additional comments (15)
src/rai_core/rai/initialization/model_initialization.py (1)
22-22: Good addition of type annotation for improved type safety.Adding the
-> BaseChatModelreturn type annotation to theget_llm_modelfunction and importing the required class improves code clarity and enables better static type checking.Also applies to: 130-130
src/rai_core/rai/aggregators/base.py (4)
15-24: Well-structured imports and type definitions.The imports are well-organized, and the use of generic type variable
Tprovides a flexible foundation for different aggregator implementations.
26-38: Good design of the BaseAggregator abstract class.The base class is well-documented and follows good OOP principles with clear separation of concerns. The use of
Generic[T]allows for type-safe implementations for different message types.
39-43: Efficient buffer management implementation.The
__call__method efficiently manages the buffer size by automatically removing the oldest message when the buffer reaches its maximum size.
44-47: Clear abstract method definition.The abstract
getmethod is well-defined with appropriate return type annotations that support both regular and multimodal message types.src/rai_core/rai/agents/ros2/__init__.py (1)
1-16: Appropriate module structure for ROS2 agent exports.The module follows standard Python packaging practices by defining
__all__to explicitly specify the public API. This provides a clean interface for users of the package.src/rai_core/rai/aggregators/__init__.py (1)
1-17: Well-structured package initialization.The module correctly exports the
BaseAggregatorclass through__all__, following Python best practices for package organization.src/rai_core/rai/agents/__init__.py (2)
15-15: Update imports to match the new architecture.The import of
BaseStateBasedAgentandStateBasedConfigfrom the new module reflects the architectural shift from the earlier factory function approach to a more modular, class-based design.
21-29: Exports updated appropriately.The
__all__list correctly exports the new classes while removing the old factory function. This ensures that the public API is up-to-date with the architectural changes.src/rai_core/rai/agents/ros2/state_based_agent.py (1)
15-17: Imports are appropriate and minimal.The imports bring in only the necessary components: the base agent class and the ROS2 connector.
src/rai_core/rai/aggregators/ros2/__init__.py (2)
15-20: Well-structured imports for aggregator classes.The imports clearly specify the four ROS2-specific aggregator classes from the local module.
22-27: Explicit exports in all list follow best practices.The
__all__list explicitly exports the four aggregator classes, defining a clear public API for the module.src/rai_core/rai/tools/ros2/generic/topics.py (1)
24-29: Updated imports to use the new conversion function.The imports are updated to include the new
encode_ros2_img_to_base64function along with the existingros2_message_to_dictfunction.src/rai_core/rai/communication/ros2/api/conversion.py (1)
29-30: Added import for preprocess_image.The import for the shared image preprocessing utility is appropriate.
src/rai_core/rai/agents/base_state_based_agent.py (1)
96-103: Callback registration leaks ifsetup_connector()failsIf any single
register_callbackraises, previously registered callbacks remain stored in the connector and_registered_callbacks, yetstop()will never be called. Wrap the loop intry/except, unregister already‑added callbacks on failure, or build the list first and register in one go.
src/rai_core/rai/aggregators/base.py
Outdated
| self._buffer.append(msg) | ||
|
|
||
| @abstractmethod | ||
| def get(self) -> HumanMessage | HumanMultimodalMessage | None: |
src/rai_core/rai/aggregators/base.py
Outdated
| return list(self._buffer) | ||
|
|
||
| def __str__(self) -> str: | ||
| return f"{self.__class__.__name__}" |
| Returns the VLM analysis of the last image in the aggregation buffer | ||
| """ | ||
|
|
||
| def __init__(self, max_size: int | None = None) -> None: |
There was a problem hiding this comment.
llm optional[BaseChatModel]
src/rai_core/rai/aggregators/base.py
Outdated
| """Returns the aggregated message""" | ||
| pass | ||
|
|
||
| def clear(self) -> None: |
| b64_images: List[str] = [encode_ros2_img_to_base64(msg) for msg in msgs] | ||
| self.clear() | ||
|
|
||
| system_prompt = "You are an expert in image analysis and your speciality is the" |
There was a problem hiding this comment.
Move to static const, allow overriding via constructor
boczekbartek
force-pushed
the
bb/feat/state_based_agent
branch
from
59520d4 to
fa820a3
Compare
boczekbartek
force-pushed
the
bb/feat/state_based_agent
branch
from
fa820a3 to
b736046
Compare
|
@CodeRabbit full review |
boczekbartek
force-pushed
the
bb/feat/state_based_agent
branch
from
cebcfb9 to
81c575a
Compare
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 2
♻️ Duplicate comments (4)
src/rai_core/rai/aggregators/__init__.py (1)
1-17: Clean package export, but check copyright year.The init file correctly exports the BaseAggregator class, but the copyright year should be updated.
-# Copyright (C) 2025 Robotec.AI +# Copyright (C) 2024 Robotec.AIsrc/rai_core/rai/agents/langchain/runnables.py (2)
128-136: In-place mutation of sharedHumanMessageobjects remains unfixed
See earlier review; the code still mutatesmessage.content, risking state corruption if the same object is reused elsewhere.
164-172: CreatingToolRunner([])still raisesValueError– original issue persists
The runnable builds aToolRunnereven whentoolsis empty, reproducing the crash highlighted in the previous review.src/rai_core/rai/aggregators/ros2/aggregators.py (1)
32-52: Duplicate-compression algorithm still loses the last duplicate block
(see previous review).The logic drops the final repetition summary and replaces the new log with
the repetition counter, discarding the actual message (BinA A A B). A
corrected version was proposed earlier and is reposted below for convenience.- prev_parsed = None - counter = 0 + prev_parsed = None + counter = 1 @@ - if parsed == prev_parsed: - counter += 1 - continue - else: - if counter != 0: - parsed = f"Log above repeated {counter} times" - buffer.append(parsed) - counter = 0 - prev_parsed = parsed + if parsed == prev_parsed: + counter += 1 + else: + # Flush previous duplicates first + if counter > 1: + buffer.append(f"Previous log repeated {counter-1} times") + buffer.append(parsed) + prev_parsed = parsed + counter = 1 + + if counter > 1: + buffer.append(f"Previous log repeated {counter-1} times") - result = f"Logs summary: {list(dict.fromkeys(buffer))}" + result = f"Logs summary: {buffer}"This preserves every unique message and accurately reports repetitions at the
end of the buffer.
🧹 Nitpick comments (13)
src/rai_core/rai/aggregators/base.py (2)
1-13: Check the copyright year.The copyright year is set to 2025, which appears to be in the future. Update it to the current year (2024) or the appropriate year when this code was written.
-# Copyright (C) 2025 Robotec.AI +# Copyright (C) 2024 Robotec.AI
23-30: Enhance documentation to describe the purpose and use cases.The current docstring explains the method behavior but doesn't describe the purpose of aggregators in the overall system.
class BaseAggregator(ABC, Generic[T]): """ Interface for aggregators. `__call__` method receives a message and appends it to the buffer. `get` method returns the aggregated message. + + Aggregators buffer messages of type T and convert them into BaseMessage instances + that can be used by LLM agents to incorporate environmental state information. + Subclasses should define how to convert buffered messages into a meaningful + representation for the agent. """docs/developer_guide.md (1)
134-135: Consider adding an explanation of "runnable" conceptThe import has been correctly updated to reflect the architectural shift from
create_state_based_agenttocreate_state_based_runnable. This aligns with the new modular framework for state-based agents introduced in the PR.Consider adding a brief note explaining what a "runnable" is in the LangChain context, as this concept might not be familiar to all developers using the framework.
src/rai_core/rai/agents/ros2/state_based_agent.py (1)
20-22: Add class docstring to explain purpose and usageThe implementation clearly defines a ROS2-specific state-based agent by overriding the
setup_connectormethod to return a ROS2Connector instance.Add a docstring to the class explaining its purpose, how it integrates with ROS2, and providing usage examples. This would be especially helpful for new users:
class ROS2StateBasedAgent(BaseStateBasedAgent): + """ + A state-based agent implementation for ROS2 environments. + + This class extends BaseStateBasedAgent to work specifically with ROS2, + setting up a ROS2Connector for communication and state aggregation. + + Example: + ```python + from rai.agents.ros2 import ROS2StateBasedAgent + from rai.aggregators.ros2 import ROS2LogsAggregator + + # Configure state aggregation + config = StateBasedConfig( + aggregators={"/rosout": [ROS2LogsAggregator()]}, + time_interval=1.0 + ) + + # Create agent + agent = ROS2StateBasedAgent(config=config, ...) + agent.run() + ``` + """ def setup_connector(self): return ROS2Connector()src/rai_core/rai/communication/ros2/api/conversion.py (1)
118-120: Usetobytes()and drop the extra copy
np.ndarray.tostring()is deprecated in NumPy ≥1.19 and will be removed.
Also, wrapping the result inbytes()above creates an unnecessary copy, unlikendarray.tobytes(), which already returnsbytes.- image_data = cv2.imencode(".png", cv_image)[1].tostring() # type: ignore + image_data = cv2.imencode(".png", cv_image)[1].tobytes() # type: ignorePlease update the earlier branches (
bytes(cv2.imencode(...)[1])) similarly for consistency.src/rai_core/rai/agents/langchain/__init__.py (2)
15-23: Heavy eager imports increase startup time & risk circular depsImporting nine sub-modules unconditionally in
__init__.pyslows downimport rai.agents.langchainand may trigger circular-import headaches (e.g. if.agentor.react_agentimport this package back).
Consider deferring the imports:from importlib import import_module import sys def __getattr__(name): mapping = { "BaseState": ".agent", "LangChainAgent": ".agent", ... } if name in mapping: module = import_module(f"{__name__}{mapping[name]}") obj = getattr(module, name) globals()[name] = obj return obj raise AttributeError(name)This keeps the public surface unchanged but avoids loading unused code.
25-36:__all__omission: newly exposed symbols must stay in syncEvery time a new submodule symbol is added you need to remember to edit
__all__.
To avoid rot, you can generate it dynamically:__all__ = [n for n in globals() if not n.startswith("_") and n[0].isupper()](or export only the mapping used in
__getattr__, if you adopt lazy loading).examples/agents/state_based.py (2)
12-13: Typo in licence boiler-plate
goveself.rninglooks like a find-replace artefact – should be “governing”.
50-53: Graceful shutdown / exception propagation
agent.run()is blocking and exceptions inside the graph will currently terminate the process without a clear message.
Consider running it in a background task and wrapping intry/finallyso thatwait_for_shutdownis always reached and connectors are closed:try: agent.run() finally: wait_for_shutdown([agent])src/rai_core/rai/agents/langchain/state_based_agent.py (2)
35-38: Fix minor typo in theStateBasedConfigdescription
"aggragator"→"aggregator".
Although harmless, public-facing docs/IDE hints inherit this typo.- description="Dict of topic : aggregator or (topic, msg_type) : aggragator" + description="Dict of topic : aggregator or (topic, msg_type) : aggregator"
172-187: Do not clear_stop_eventafter shutdown
stop()sets_stop_eventto notify all threads that shutdown is in
progress. Clearing it at the end of the method re-opens a tiny time window in
which the agent instance appears “running” again, which can lead to surprises
if the same object is reused or inspected from another thread.- self._stop_event.clear()Recommendation: remove the line unless you have a documented requirement to
restart the very same agent instance.src/rai_core/rai/aggregators/ros2/aggregators.py (2)
165-169: Misreported image counts in VLM diff outputAfter calling
get_key_elements,len(b64_images)refers to the selected
subset, not the total number of frames analysed. The summary therefore prints
the same number twice.- content=f"Result of the analysis of the {len(b64_images)} keyframes selected from {len(b64_images)} last images:\n{response}" + original_count = len(all_images) # capture before subsampling + keyframes = self.get_key_elements(all_images) + ... + content=( + f"Result of the analysis of the {len(keyframes)} keyframes " + f"selected from {original_count} last images:\n{response}" + )Capturing the original buffer length before subsampling (and using distinct
variable names) yields a correct, reader-friendly message.
113-118: Minor doc-string typo“midden” → “middle”.
- aggregation buffer: 1st, midden, last + aggregation buffer: 1st, middle, last
🛑 Comments failed to post (2)
src/rai_core/rai/communication/ros2/api/conversion.py (1)
95-106: 🛠️ Refactor suggestion
⚠️ Potential issueGrayscale handling & channel-count detection can silently mis-route images
cv_image.shape[-1]assumes a 3-D ndarray.
For mono images returned byCvBridge,shape == (h, w)(2-D), soshape[-1]equalsw, not1.
The current logic therefore falls through to the BGR path, giving an incorrect colour conversion and distorted PNG.- if cv_image.shape[-1] == 4: + # BGRA + if len(cv_image.shape) == 3 and cv_image.shape[-1] == 4: ... - elif cv_image.shape[-1] == 1: + # MONO + elif len(cv_image.shape) == 2 or (len(cv_image.shape) == 3 and cv_image.shape[-1] == 1):Adding the
len(cv_image.shape)guard ensures both(h,w)and(h,w,1)grayscale inputs select the intended branch.
Please add a test with a MONO8sensor_msgs.msg.Imageto verify.Committable suggestion skipped: line range outside the PR's diff.
src/rai_core/rai/agents/langchain/state_based_agent.py (1)
86-118: 🛠️ Refactor suggestion
Guard
_aggregation_resultswith a lock to avoid race conditions
_aggregation_resultsis written in the background aggregation thread and may be
read concurrently by:
- The runnable graph (through
self.get_state()).- Any external callers that query
get_state().Without synchronisation, a read during a concurrent write can raise
RuntimeError: dictionary changed size during iterationor return a partially
updated state.Suggested minimal fix:
@@ - self._aggregation_results: Dict[str, HumanMessage | HumanMultimodalMessage] = ( - dict() - ) + self._aggregation_results: Dict[str, HumanMessage | HumanMultimodalMessage] = {} + self._state_lock = threading.Lock() @@ def get_state(self) -> Dict[str, HumanMessage | HumanMultimodalMessage]: """Returns output for all aggregators""" - return self._aggregation_results + with self._state_lock: + # Return a shallow copy to prevent callers from mutating internal state + return dict(self._aggregation_results) @@ - self._aggregation_results[source] = output + with self._state_lock: + self._aggregation_results[source] = outputThis keeps the change local and inexpensive while guaranteeing
thread-safety.
If higher throughput is required later, consider switching to a
collections.OrderedDictguarded bythreading.RLockor an immutable
snapshotting strategy.Also applies to: 168-171
Purpose
Proposed Changes
StateBaseAgentthat can expose the state to the LLM agent and collect stateusing aggragation functions.
BaseStateAggregatorTODOs:
Issues
Testing
source setup_shell.sh python examples/agents/state_based.pyros2 topic pub /from_human rai_interfaces/msg/HRIMessage "{'text': 'Is nav2 stack working correctly?'}" --onceSummary by CodeRabbit
New Features
Bug Fixes
Documentation
Chores