add logging to agentchat

This commit is contained in:
Victor Dibia 2024-10-02 17:25:14 -07:00
parent 9ba14ee15b
commit b0b0825c1d
7 changed files with 144 additions and 25 deletions

View File

@ -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

View File

@ -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()

View File

@ -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
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.
self._logger.info(ContentPublishEvent(agent_message=response))
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")
# Forward the feedback to the agent.
self._logger.info(ContentPublishEvent(agent_message=feedback, source=self._tool_agent_id.type))
response = await self._agent.on_messages([feedback], ctx.cancellation_token)
# Publish the response.

View File

@ -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")

View File

@ -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):

View File

@ -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)

View File

@ -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)