mirror of
https://github.com/microsoft/autogen.git
synced 2025-06-26 22:30:10 +00:00
db125fbd2d
## Why are these changes needed? I added `created_at` to both BaseChatMessage and BaseAgentEvent classes that store the time these Pydantic model instances are generated. And then users will be able to use `created_at` to build up a customized external persisting state management layer for their case. ## Related issue number https://github.com/microsoft/autogen/discussions/6169#discussioncomment-13151540 ## Checks - [x] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> to build and test documentation locally. - [x] I've added tests (if relevant) corresponding to the changes introduced in this PR. - [x] I've made sure all auto checks have passed. --------- Co-authored-by: Jack Gerrits <jackgerrits@users.noreply.github.com> Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
import json
|
|
import logging
|
|
import sys
|
|
from datetime import datetime
|
|
from typing import Sequence
|
|
|
|
from autogen_agentchat.base._task import TaskResult
|
|
from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, BaseTextChatMessage
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class FileLogHandler(logging.Handler):
|
|
def __init__(self, filename: str) -> None:
|
|
super().__init__()
|
|
self.filename = filename
|
|
self.file_handler = logging.FileHandler(filename)
|
|
|
|
def emit(self, record: logging.LogRecord) -> None:
|
|
ts = datetime.fromtimestamp(record.created).isoformat()
|
|
if isinstance(record.msg, BaseModel):
|
|
record.msg = json.dumps(
|
|
{
|
|
"timestamp": ts,
|
|
"message": record.msg.model_dump_json(indent=2),
|
|
"type": record.msg.__class__.__name__,
|
|
},
|
|
)
|
|
self.file_handler.emit(record)
|
|
|
|
|
|
class ConsoleLogHandler(logging.Handler):
|
|
def emit(self, record: logging.LogRecord) -> None:
|
|
ts = datetime.fromtimestamp(record.created).isoformat()
|
|
if isinstance(record.msg, BaseModel):
|
|
record.msg = json.dumps(
|
|
{
|
|
"timestamp": ts,
|
|
"message": record.msg.model_dump_json(indent=2),
|
|
"type": record.msg.__class__.__name__,
|
|
},
|
|
)
|
|
sys.stdout.write(f"{record.msg}\n")
|
|
|
|
|
|
def compare_messages(
|
|
msg1: BaseAgentEvent | BaseChatMessage | BaseTextChatMessage,
|
|
msg2: BaseAgentEvent | BaseChatMessage | BaseTextChatMessage,
|
|
) -> bool:
|
|
if isinstance(msg1, BaseTextChatMessage) and isinstance(msg2, BaseTextChatMessage):
|
|
if msg1.content != msg2.content:
|
|
return False
|
|
return (
|
|
(msg1.source == msg2.source) and (msg1.models_usage == msg2.models_usage) and (msg1.metadata == msg2.metadata)
|
|
)
|
|
|
|
|
|
def compare_message_lists(
|
|
msgs1: Sequence[BaseAgentEvent | BaseChatMessage],
|
|
msgs2: Sequence[BaseAgentEvent | BaseChatMessage],
|
|
) -> bool:
|
|
if len(msgs1) != len(msgs2):
|
|
return False
|
|
for i in range(len(msgs1)):
|
|
if not compare_messages(msgs1[i], msgs2[i]):
|
|
return False
|
|
return True
|
|
|
|
|
|
def compare_task_results(
|
|
res1: TaskResult,
|
|
res2: TaskResult,
|
|
) -> bool:
|
|
if res1.stop_reason != res2.stop_reason:
|
|
return False
|
|
return compare_message_lists(res1.messages, res2.messages)
|