support multiple callback fn; improve examples of how to add callback from script (#38)

Author: xingyaoww

examples/hello_world.py

@@ -4,10 +4,13 @@
 
 from openhands.core import (
     LLM,
+    ActionBase,
     CodeActAgent,
     Conversation,
+    ConversationEventType,
     LLMConfig,
     Message,
+    ObservationBase,
     TextContent,
     Tool,
     get_logger,
@@ -42,7 +45,19 @@
 
 # Agent
 agent = CodeActAgent(llm=llm, tools=tools)
-conversation = Conversation(agent=agent)
+
+llm_messages = []  # collect raw LLM messages
+def conversation_callback(event: ConversationEventType):
+    # print all the actions
+    if isinstance(event, ActionBase):
+        logger.info(f"Found a conversation action: {event}")
+    elif isinstance(event, ObservationBase):
+        logger.info(f"Found a conversation observation: {event}")
+    elif isinstance(event, Message):
+        logger.info(f"Found a conversation message: {str(event)[:200]}...")
+        llm_messages.append(event.model_dump())
+
+conversation = Conversation(agent=agent, callbacks=[conversation_callback])
 
 conversation.send_message(
     message=Message(
@@ -51,3 +66,8 @@
     )
 )
 conversation.run()
+
+print("="*100)
+print("Conversation finished. Got the following LLM messages:")
+for i, message in enumerate(llm_messages):
+    print(f"Message {i}: {str(message)[:200]}")

openhands/core/__init__.py

@@ -2,7 +2,7 @@
 
 from .agent import AgentBase, CodeActAgent
 from .config import LLMConfig, MCPConfig
-from .conversation import Conversation
+from .conversation import Conversation, ConversationCallbackType, ConversationEventType
 from .llm import LLM, ImageContent, Message, TextContent
 from .logger import get_logger
 from .tool import ActionBase, ObservationBase, Tool
@@ -27,5 +27,7 @@
     "MCPConfig",
     "get_logger",
     "Conversation",
+    "ConversationCallbackType",
+    "ConversationEventType",
     "__version__",
 ]

openhands/core/agent/base.py

@@ -3,7 +3,7 @@
 
 from openhands.core.context.env_context import EnvContext
 from openhands.core.conversation import ConversationCallbackType, ConversationState
-from openhands.core.llm import LLM
+from openhands.core.llm import LLM, Message
 from openhands.core.logger import get_logger
 from openhands.core.tool import Tool
 
@@ -60,6 +60,7 @@ def env_context(self) -> EnvContext | None:
     def init_state(
         self,
         state: ConversationState,
+        initial_user_message: Message | None = None,
         on_event: ConversationCallbackType | None = None,
     ) -> ConversationState:
         """Initialize the empty conversation state to prepare the agent for user messages.

openhands/core/agent/codeact_agent/codeact_agent.py

@@ -77,16 +77,25 @@ def __init__(
     def init_state(
         self,
         state: ConversationState,
+        initial_user_message: Message | None = None,
         on_event: ConversationCallbackType | None = None,
     ) -> ConversationState:
         # TODO(openhands): we should add test to test this init_state will actually modify state in-place
         messages = state.history.messages
         if len(messages) == 0:
+            # Prepare system message
             sys_msg = Message(role="system", content=[self.system_message])
             messages.append(sys_msg)
             if on_event:
                 on_event(sys_msg)
-            content = state.history.messages[-1].content
+            if initial_user_message is None:
+                raise ValueError("initial_user_message must be provided in init_state for CodeActAgent")
+            
+            # Prepare user message
+            content = initial_user_message.content
+            # TODO: think about this - we might want to handle this outside Agent but inside Conversation (e.g., in send_messages)
+            # downside of handling them inside Conversation would be: conversation don't have access
+            # to *any* action execution runtime information
             if self.env_context:
                 initial_env_context: list[TextContent] = self.env_context.render(self.prompt_manager)
                 content += initial_env_context

openhands/core/conversation/__init__.py

@@ -1,7 +1,7 @@
 from .conversation import Conversation
 from .state import ConversationState
-from .types import ConversationCallbackType
+from .types import ConversationCallbackType, ConversationEventType
 from .visualizer import ConversationVisualizer
 
 
-__all__ = ["Conversation", "ConversationState", "ConversationCallbackType", "ConversationVisualizer"]
+__all__ = ["Conversation", "ConversationState", "ConversationCallbackType", "ConversationEventType", "ConversationVisualizer"]

openhands/core/conversation/conversation.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterable
 
 
 if TYPE_CHECKING:
@@ -16,30 +16,52 @@
 logger = get_logger(__name__)
 
 
+def compose_callbacks(callbacks: Iterable[ConversationCallbackType]) -> ConversationCallbackType:
+    def composed(event) -> None:
+        for cb in callbacks:
+            if cb:
+                cb(event)
+    return composed
+
 class Conversation:
-    def __init__(self, agent: "AgentBase", on_event: ConversationCallbackType | None = None, max_iteration_per_run: int = 500):
+    def __init__(
+        self,
+        agent: "AgentBase",
+        callbacks: list[ConversationCallbackType] | None = None,
+        max_iteration_per_run: int = 500,
+    ):
+        """Initialize the conversation."""
         self._visualizer = ConversationVisualizer()
-        self._on_event: ConversationCallbackType = on_event or self._visualizer.on_event
+        # Compose multiple callbacks if a list is provided
+        self._on_event = compose_callbacks(
+            [self._visualizer.on_event] + (callbacks if callbacks else [])
+        )
         self.max_iteration_per_run = max_iteration_per_run
 
         self.agent = agent
+        self._agent_initialized = False
 
         # Guarding the conversation state to prevent multiple
         # writers modify it at the same time
         self._lock = RLock()
         self.state = ConversationState()
 
-        with self._lock:
-            # will modify self.state in place
-            self.state = self.agent.init_state(self.state, on_event=self._on_event)
-
     def send_message(self, message: Message) -> None:
         """Sending messages to the agent."""
         with self._lock:
-            messages = self.state.history.messages
-            messages.append(message)
-            if self._on_event:
-                self._on_event(message)
+            if not self._agent_initialized:
+                # Prepare initial state
+                self.state = self.agent.init_state(
+                    self.state,
+                    initial_user_message=message,
+                    on_event=self._on_event,
+                )
+                self._agent_initialized = True
+            else:
+                messages = self.state.history.messages
+                messages.append(message)
+                if self._on_event:
+                    self._on_event(message)
 
     def run(self) -> None:
         """Runs the conversation until the agent finishes."""

openhands/core/conversation/types.py

@@ -4,4 +4,5 @@
 from openhands.core.tool import ActionBase, ObservationBase
 
 
-ConversationCallbackType = Callable[[Message | ActionBase | ObservationBase], None]
+ConversationEventType = Message | ActionBase | ObservationBase
+ConversationCallbackType = Callable[[ConversationEventType], None]

openhands/core/llm/llm.py

@@ -676,6 +676,7 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return str(self)
 
+    # TODO: we should ideally format this into a `to_litellm_message` for `Message` class`
     def format_messages_for_llm(self, messages: Message | list[Message]) -> list[dict]:
         if isinstance(messages, Message):
             messages = [messages]