feat: add StateBasedAgent and rai.aggregators

Author: boczekbartek

docs/developer_guide.md

@@ -131,7 +131,7 @@ Choose the conversational_agent when beginning with RAI, and consider upgrading
 from myrobot import robot
 
 from rai.agents.conversational_agent import create_conversational_agent
-from rai.agents.state_based import create_state_based_agent
+from rai.agents.langchain import create_state_based_runnable
 from rai import get_llm_model
 
 SYSTEM_PROMPT = "You are a robot with interfaces..."
@@ -143,7 +143,7 @@ tools = [pick_up_object, scan_object, SayTool(robot=robot)]
 conversational_agent = create_conversational_agent(
     llm=llm, tools=tools, system_prompt=SYSTEM_PROMPT
 )
-state_based_agent = create_state_based_agent(
+state_based_agent = create_state_based_runnable(
     llm=llm, state_retriever=state_retriever, tools=tools
 )
 

examples/agents/state_based.py

@@ -0,0 +1,53 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language goveself.rning permissions and
+# limitations under the License.
+
+from rai.agents import wait_for_shutdown
+from rai.agents.langchain import StateBasedConfig
+from rai.agents.ros2 import ROS2StateBasedAgent
+from rai.aggregators.ros2 import (
+    ROS2ImgVLMDiffAggregator,
+    ROS2LogsAggregator,
+)
+from rai.communication.ros2 import (
+    ROS2Connector,
+    ROS2Context,
+    ROS2HRIConnector,
+)
+from rai.tools.ros2.generic.toolkit import ROS2Toolkit
+
+
+@ROS2Context()
+def main():
+    hri_connector = ROS2HRIConnector()
+    ros2_connector = ROS2Connector()
+
+    config = StateBasedConfig(
+        aggregators={
+            "/camera/camera/color/image_raw": [ROS2ImgVLMDiffAggregator()],
+            "/rosout": [ROS2LogsAggregator()],
+        }
+    )
+
+    agent = ROS2StateBasedAgent(
+        config=config,
+        target_connectors={"to_human": hri_connector},
+        tools=ROS2Toolkit(connector=ros2_connector).get_tools(),
+    )
+    agent.subscribe_source("/from_human", hri_connector)
+    agent.run()
+    wait_for_shutdown([agent])
+
+
+if __name__ == "__main__":
+    main()

src/rai_core/rai/agents/__init__.py

@@ -16,7 +16,6 @@
 from rai.agents.conversational_agent import create_conversational_agent
 from rai.agents.langchain.react_agent import ReActAgent
 from rai.agents.runner import AgentRunner, wait_for_shutdown
-from rai.agents.state_based import create_state_based_agent
 from rai.agents.tool_runner import ToolRunner
 
 __all__ = [
@@ -25,6 +24,5 @@
     "ReActAgent",
     "ToolRunner",
     "create_conversational_agent",
-    "create_state_based_agent",
     "wait_for_shutdown",
 ]

src/rai_core/rai/agents/langchain/__init__.py

@@ -12,16 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .agent import BaseState, LangChainAgent
+from .agent import BaseState, LangChainAgent, newMessageBehaviorType
 from .callback import HRICallbackHandler
 from .react_agent import ReActAgent
-from .runnables import ReActAgentState, create_react_runnable
+from .runnables import (
+    ReActAgentState,
+    create_react_runnable,
+    create_state_based_runnable,
+)
+from .state_based_agent import BaseStateBasedAgent, StateBasedConfig
 
 __all__ = [
     "BaseState",
+    "BaseStateBasedAgent",
     "HRICallbackHandler",
     "LangChainAgent",
     "ReActAgent",
     "ReActAgentState",
+    "StateBasedConfig",
     "create_react_runnable",
+    "create_state_based_runnable",
+    "newMessageBehaviorType",
 ]

src/rai_core/rai/agents/langchain/agent.py

@@ -22,12 +22,13 @@
 from langchain_core.messages import BaseMessage
 from langchain_core.runnables import Runnable
 
-from rai.agents.base import BaseAgent
-from rai.agents.langchain import HRICallbackHandler
-from rai.agents.langchain.runnables import ReActAgentState
 from rai.communication.hri_connector import HRIConnector, HRIMessage
 from rai.initialization import get_tracing_callbacks
 
+from ..base import BaseAgent
+from .callback import HRICallbackHandler
+from .runnables import ReActAgentState
+
 
 class BaseState(TypedDict):
     messages: List[BaseMessage]

src/rai_core/rai/agents/langchain/react_agent.py

@@ -17,11 +17,11 @@
 from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import BaseTool
 
-from rai.agents.langchain import create_react_runnable
-from rai.agents.langchain.agent import LangChainAgent
-from rai.agents.langchain.runnables import ReActAgentState
 from rai.communication.hri_connector import HRIConnector, HRIMessage
 
+from .agent import LangChainAgent
+from .runnables import ReActAgentState, create_react_runnable
+
 
 class ReActAgent(LangChainAgent):
     def __init__(

src/rai_core/rai/agents/langchain/runnables.py

@@ -31,10 +31,11 @@
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt.tool_node import tools_condition
 
-from rai.agents.tool_runner import ToolRunner
 from rai.initialization import get_llm_model
 from rai.messages import HumanMultimodalMessage
 
+from ..tool_runner import ToolRunner
+
 
 class ReActAgentState(TypedDict):
     """State type for the react agent.

src/rai_core/rai/agents/base_state_based_agent.py renamed to src/rai_core/rai/agents/langchain/state_based_agent.py

@@ -22,16 +22,14 @@
 from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel, ConfigDict, Field
-from rclpy.callback_groups import ReentrantCallbackGroup
-from rclpy.subscription import Subscription
 
-from rai.agents.langchain import create_state_based_runnable
 from rai.aggregators import BaseAggregator
 from rai.communication.base_connector import BaseConnector
 from rai.communication.hri_connector import HRIConnector, HRIMessage
 from rai.messages.multimodal import HumanMultimodalMessage
 
-from .langchain import ReActAgent, ReActAgentState, create_state_based_runnable
+from .agent import LangChainAgent, newMessageBehaviorType
+from .runnables import ReActAgentState, create_state_based_runnable
 
 
 class StateBasedConfig(BaseModel):
@@ -44,7 +42,7 @@ class StateBasedConfig(BaseModel):
     )
 
 
-class BaseStateBasedAgent(ReActAgent, ABC):
+class BaseStateBasedAgent(LangChainAgent, ABC):
     """
     Agent that runs aggregators (config.aggregators) every config.time_interval seconds.
     Aggregators are registered to their sources using
@@ -59,12 +57,14 @@ class BaseStateBasedAgent(ReActAgent, ABC):
 
     def __init__(
         self,
-        connectors: dict[str, HRIConnector[HRIMessage]],
         config: StateBasedConfig,
+        target_connectors: Dict[str, HRIConnector[HRIMessage]],
         llm: Optional[BaseChatModel] = None,
         tools: Optional[List[BaseTool]] = None,
         state: Optional[ReActAgentState] = None,
         system_prompt: Optional[str] = None,
+        new_message_behavior: newMessageBehaviorType = "interrupt_keep_last",
+        max_size: int = 100,
     ):
         runnable = create_state_based_runnable(
             llm=llm,
@@ -73,13 +73,14 @@ def __init__(
             state_retriever=self.get_state,
         )
         super().__init__(
-            connectors, llm, tools, state, system_prompt, runnable=runnable
+            target_connectors=target_connectors,
+            runnable=runnable,
+            state=state,
+            new_message_behavior=new_message_behavior,
+            max_size=max_size,
         )
         self.config = config
 
-        self._callback_group = ReentrantCallbackGroup()
-        self._subscriptions: Dict[str, Subscription] = dict()
-
         self._aggregation_results: Dict[str, HumanMessage | HumanMultimodalMessage] = (
             dict()
         )

src/rai_core/rai/agents/ros2/state_based_agent.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from rai.agents import BaseStateBasedAgent
 from rai.communication.ros2 import ROS2Connector
 
+from ..langchain import BaseStateBasedAgent
+
 
 class ROS2StateBasedAgent(BaseStateBasedAgent):
     def setup_connector(self):