add logging to agentchat

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_base_team.py

@@ -1,7 +1,16 @@
 from dataclasses import dataclass
+import logging
 from typing import List, Protocol
 
 from autogen_agentchat.agents._base_chat_agent import ChatMessage
+from autogen_core.application.logging import EVENT_LOGGER_NAME
+
+from ._utils import AgentChatLogHandler
+
+logger = logging.getLogger(EVENT_LOGGER_NAME + ".agentchatchat")
+logger.setLevel(logging.INFO)
+log_handler = AgentChatLogHandler(filename="log.jsonl")
+logger.handlers = [log_handler]
 
 
 @dataclass

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_utils.py

@@ -0,0 +1,89 @@
+from datetime import datetime
+import json
+import logging
+import sys
+from typing import Optional, Union, List, Dict, Any, Sequence
+from dataclasses import asdict, is_dataclass
+
+from .group_chat._events import ContentPublishEvent
+from ..agents import ChatMessage, TextMessage, MultiModalMessage, ToolCallMessage, ToolCallResultMessage, StopMessage
+from autogen_core.components import FunctionCall, Image
+from autogen_core.components.models import FunctionExecutionResult
+
+ContentType = Union[str, List[Union[str, Image]],
+                    List[FunctionCall], List[FunctionExecutionResult]]
+
+
+class AgentChatLogHandler(logging.Handler):
+    def __init__(self, filename: Optional[str] = None) -> None:
+        super().__init__()
+        self.filename = filename
+        self.file_handler: Optional[logging.FileHandler] = None
+        if filename:
+            self.file_handler = logging.FileHandler(filename)
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            ts = datetime.fromtimestamp(record.created).isoformat()
+            if isinstance(record.msg, ContentPublishEvent):
+                console_message = (
+                    f"\n{'-'*75} \n"
+                    f"\033[91m[{ts}], {record.msg.agent_message.source}:\033[0m\n"
+                    f"\n{self.serialize_content(record.msg.agent_message.content)}"
+                )
+                # print , flush true
+                sys.stdout.write(console_message)
+
+                log_entry = json.dumps(
+                    {
+                        "timestamp": ts,
+                        "source": record.msg.agent_message.source,
+                        "message": self.serialize_content(record.msg.agent_message.content),
+                        "type": "OrchestrationEvent",
+                    },
+                    default=self.json_serializer,
+                )
+
+                if self.file_handler:
+                    file_record = logging.LogRecord(
+                        name=record.name,
+                        level=record.levelno,
+                        pathname=record.pathname,
+                        lineno=record.lineno,
+                        msg=log_entry,
+                        args=(),
+                        exc_info=record.exc_info,
+                    )
+                    self.file_handler.emit(file_record)
+                else:
+                    sys.stderr.write(log_entry)
+        except Exception:
+            self.handleError(record)
+
+    def serialize_content(
+        self, content: Union[ContentType, Sequence[ChatMessage], ChatMessage]
+    ) -> Union[List[Any], Dict[str, Any], str]:
+        if isinstance(content, (str, list)):
+            return content
+        elif isinstance(content, (TextMessage, MultiModalMessage, ToolCallMessage, ToolCallResultMessage, StopMessage)):
+            return asdict(content)
+        elif isinstance(content, Image):
+            return {"type": "image", "data": content.data_uri}
+        elif isinstance(content, FunctionCall):
+            return {"type": "function_call", "name": content.name, "arguments": content.arguments}
+        elif isinstance(content, FunctionExecutionResult):
+            return {"type": "function_execution_result", "content": content.content}
+        return str(content)
+
+    @staticmethod
+    def json_serializer(obj: Any) -> Any:
+        if is_dataclass(obj) and not isinstance(obj, type):
+            return asdict(obj)
+        elif isinstance(obj, type):
+            return str(obj)
+        return str(obj)
+
+    def close(self) -> None:
+        if self.file_handler:
+            self.file_handler.close()
+        super().close()

python/packages/autogen-agentchat/src/autogen_agentchat/teams/group_chat/_base_chat_agent_container.py

@@ -1,11 +1,13 @@
+import logging
 import asyncio
-import sys
 from typing import List
 
+
 from autogen_core.base import AgentId, AgentType, MessageContext
 from autogen_core.components import DefaultTopicId, event
 from autogen_core.components.models import FunctionExecutionResult
 from autogen_core.components.tool_agent import ToolException
+from autogen_core.application.logging import EVENT_LOGGER_NAME
 
 from ...agents import BaseChatAgent, MultiModalMessage, StopMessage, TextMessage, ToolCallMessage, ToolCallResultMessage
 from ._events import ContentPublishEvent, ContentRequestEvent
@@ -29,6 +31,7 @@ class BaseChatAgentContainer(SequentialRoutedAgent):
         self._agent = agent
         self._message_buffer: List[TextMessage | MultiModalMessage | StopMessage] = []
         self._tool_agent_id = AgentId(type=tool_agent_type, key=self.id.key)
+        self._logger = self.logger = logging.getLogger(EVENT_LOGGER_NAME + f".agentchatchat")
 
     @event
     async def handle_content_publish(self, message: ContentPublishEvent, ctx: MessageContext) -> None:
@@ -48,9 +51,8 @@ class BaseChatAgentContainer(SequentialRoutedAgent):
 
         # Handle tool calls.
         while isinstance(response, ToolCallMessage):
-            # TODO: use logging instead of print
+            self._logger.info(ContentPublishEvent(agent_message=response))
-            sys.stdout.write(f"{'-'*80}\n{self._agent.name}:\n{response.content}\n")
+
-            # Execute functions called by the model by sending messages to tool agent.
             results: List[FunctionExecutionResult | BaseException] = await asyncio.gather(
                 *[
                     self.send_message(
@@ -73,8 +75,7 @@ class BaseChatAgentContainer(SequentialRoutedAgent):
             # Create a new tool call result message.
             feedback = ToolCallResultMessage(content=function_results, source=self._tool_agent_id.type)
             # TODO: use logging instead of print
-            sys.stdout.write(f"{'-'*80}\n{self._tool_agent_id.type}:\n{feedback.content}\n")
+            self._logger.info(ContentPublishEvent(agent_message=feedback, source=self._tool_agent_id.type))
-            # Forward the feedback to the agent.
             response = await self._agent.on_messages([feedback], ctx.cancellation_token)
 
         # Publish the response.

python/packages/autogen-agentchat/src/autogen_agentchat/teams/group_chat/_base_group_chat_manager.py

@@ -1,12 +1,13 @@
-import sys
+import logging
 from typing import List
 
 from autogen_core.base import MessageContext, TopicId
 from autogen_core.components import event
 
-from ...agents import MultiModalMessage, StopMessage, TextMessage
+from ...agents import StopMessage, TextMessage, ChatMessage
 from ._events import ContentPublishEvent, ContentRequestEvent
 from ._sequential_routed_agent import SequentialRoutedAgent
+from autogen_core.application.logging import EVENT_LOGGER_NAME
 
 
 class BaseGroupChatManager(SequentialRoutedAgent):
@@ -48,7 +49,8 @@ class BaseGroupChatManager(SequentialRoutedAgent):
             raise ValueError("The group topic type must not be the same as the parent topic type.")
         self._participant_topic_types = participant_topic_types
         self._participant_descriptions = participant_descriptions
-        self._message_thread: List[TextMessage | MultiModalMessage | StopMessage] = []
+        self._message_thread: List[ChatMessage] = []
+        self._logger = self.logger = logging.getLogger(EVENT_LOGGER_NAME + ".agentchatchat")
 
     @event
     async def handle_content_publish(self, message: ContentPublishEvent, ctx: MessageContext) -> None:
@@ -62,7 +64,8 @@ class BaseGroupChatManager(SequentialRoutedAgent):
         group_chat_topic_id = TopicId(type=self._group_topic_type, source=ctx.topic_id.source)
 
         # TODO: use something else other than print.
-        sys.stdout.write(f"{'-'*80}\n{message.agent_message.source}:\n{message.agent_message.content}\n")
+
+        self._logger.info(ContentPublishEvent(agent_message=message.agent_message))
 
         # Process event from parent.
         if ctx.topic_id.type == self._parent_topic_type:
@@ -105,7 +108,7 @@ class BaseGroupChatManager(SequentialRoutedAgent):
         participant_topic_id = TopicId(type=speaker_topic_type, source=ctx.topic_id.source)
         await self.publish_message(ContentRequestEvent(), topic_id=participant_topic_id)
 
-    async def select_speaker(self, thread: List[TextMessage | MultiModalMessage | StopMessage]) -> str:
+    async def select_speaker(self, thread: List[ChatMessage]) -> str:
         """Select a speaker from the participants and return the
         topic type of the selected speaker."""
         raise NotImplementedError("Method not implemented")

python/packages/autogen-agentchat/src/autogen_agentchat/teams/group_chat/_events.py

@@ -1,6 +1,7 @@
+from typing import Optional
 from pydantic import BaseModel
 
-from ...agents import MultiModalMessage, StopMessage, TextMessage
+from ...agents import ChatMessage
 
 
 class ContentPublishEvent(BaseModel):
@@ -9,8 +10,9 @@ class ContentPublishEvent(BaseModel):
     content of the event.
     """
 
-    agent_message: TextMessage | MultiModalMessage | StopMessage
+    agent_message: ChatMessage
     """The message published by the agent."""
+    source: Optional[str] = None
 
 
 class ContentRequestEvent(BaseModel):

python/packages/autogen-agentchat/src/autogen_agentchat/teams/group_chat/_round_robin_group_chat.py

@@ -8,6 +8,7 @@ from autogen_core.components import ClosureAgent, TypeSubscription
 from autogen_core.components.tool_agent import ToolAgent
 from autogen_core.components.tools import Tool
 
+
 from ...agents import BaseChatAgent, TextMessage
 from .._base_team import BaseTeam, TeamRunResult
 from ._base_chat_agent_container import BaseChatAgentContainer
@@ -54,7 +55,12 @@ class RoundRobinGroupChat(BaseTeam):
 
     """
 
-    def __init__(self, participants: List[BaseChatAgent], *, tools: List[Tool] | None = None):
+    def __init__(
+        self,
+        participants: List[BaseChatAgent],
+        *,
+        tools: List[Tool] | None = None
+    ):
         if len(participants) == 0:
             raise ValueError("At least one participant is required.")
         if len(participants) != len(set(participant.name for participant in participants)):
@@ -69,7 +75,8 @@ class RoundRobinGroupChat(BaseTeam):
         def _factory() -> BaseChatAgentContainer:
             id = AgentInstantiationContext.current_agent_id()
             assert id == AgentId(type=agent.name, key=self._team_id)
-            container = BaseChatAgentContainer(parent_topic_type, agent, tool_agent_type)
+            container = BaseChatAgentContainer(
+                parent_topic_type, agent, tool_agent_type)
             assert container.id == id
             return container
 
@@ -88,7 +95,8 @@ class RoundRobinGroupChat(BaseTeam):
 
         # Register the tool agent.
         tool_agent_type = await ToolAgent.register(
-            runtime, "tool_agent", lambda: ToolAgent("Tool agent for round-robin group chat", self._tools)
+            runtime, "tool_agent", lambda: ToolAgent(
+                "Tool agent for round-robin group chat", self._tools)
         )
         # No subscriptions are needed for the tool agent, which will be called via direct messages.
 
@@ -101,7 +109,8 @@ class RoundRobinGroupChat(BaseTeam):
             topic_type = participant.name
             # Register the participant factory.
             await BaseChatAgentContainer.register(
-                runtime, type=agent_type, factory=self._create_factory(group_topic_type, participant, tool_agent_type)
+                runtime, type=agent_type, factory=self._create_factory(
+                    group_topic_type, participant, tool_agent_type)
             )
             # Add subscriptions for the participant.
             await runtime.add_subscription(TypeSubscription(topic_type=topic_type, agent_type=agent_type))
@@ -123,13 +132,16 @@ class RoundRobinGroupChat(BaseTeam):
         )
         # Add subscriptions for the group chat manager.
         await runtime.add_subscription(
-            TypeSubscription(topic_type=group_chat_manager_topic_type, agent_type=group_chat_manager_agent_type.type)
+            TypeSubscription(topic_type=group_chat_manager_topic_type,
+                             agent_type=group_chat_manager_agent_type.type)
         )
         await runtime.add_subscription(
-            TypeSubscription(topic_type=group_topic_type, agent_type=group_chat_manager_agent_type.type)
+            TypeSubscription(topic_type=group_topic_type,
+                             agent_type=group_chat_manager_agent_type.type)
         )
         await runtime.add_subscription(
-            TypeSubscription(topic_type=team_topic_type, agent_type=group_chat_manager_agent_type.type)
+            TypeSubscription(topic_type=team_topic_type,
+                             agent_type=group_chat_manager_agent_type.type)
         )
 
         group_chat_messages: List[ChatMessage] = []
@@ -144,7 +156,8 @@ class RoundRobinGroupChat(BaseTeam):
             type="collect_group_chat_messages",
             closure=collect_group_chat_messages,
             subscriptions=lambda: [
-                TypeSubscription(topic_type=group_topic_type, agent_type="collect_group_chat_messages")
+                TypeSubscription(topic_type=group_topic_type,
+                                 agent_type="collect_group_chat_messages")
             ],
         )
 
@@ -153,9 +166,11 @@ class RoundRobinGroupChat(BaseTeam):
 
         # Run the team by publishing the task to the team topic and then requesting the result.
         team_topic_id = TopicId(type=team_topic_type, source=self._team_id)
-        group_chat_manager_topic_id = TopicId(type=group_chat_manager_topic_type, source=self._team_id)
+        group_chat_manager_topic_id = TopicId(
+            type=group_chat_manager_topic_type, source=self._team_id)
         await runtime.publish_message(
-            ContentPublishEvent(agent_message=TextMessage(content=task, source="user")),
+            ContentPublishEvent(agent_message=TextMessage(
+                content=task, source="user")),
             topic_id=team_topic_id,
         )
         await runtime.publish_message(ContentRequestEvent(), topic_id=group_chat_manager_topic_id)

python/packages/autogen-agentchat/src/autogen_agentchat/teams/group_chat/_round_robin_group_chat_manager.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from ...agents import MultiModalMessage, StopMessage, TextMessage
+from ...agents import ChatMessage
 from ._base_group_chat_manager import BaseGroupChatManager
 
 
@@ -22,7 +22,7 @@ class RoundRobinGroupChatManager(BaseGroupChatManager):
         )
         self._next_speaker_index = 0
 
-    async def select_speaker(self, thread: List[TextMessage | MultiModalMessage | StopMessage]) -> str:
+    async def select_speaker(self, thread: List[ChatMessage]) -> str:
         """Select a speaker from the participants in a round-robin fashion."""
         current_speaker_index = self._next_speaker_index
         self._next_speaker_index = (current_speaker_index + 1) % len(self._participant_topic_types)