From 9df928b73ea4a990d74ff60d9cfac8ce2fcb66ec Mon Sep 17 00:00:00 2001
From: gagb <gagb@users.noreply.github.com>
Date: Tue, 2 Jul 2024 18:09:45 -0700
Subject: [PATCH] Improve printing further (#172)

* Improve pretty printing

* Improve printing further
---
 .../HumanEval/Templates/TeamOne/scenario.py    |  5 ++++-
 .../src/team_one/agents/orchestrator.py        | 18 ++++--------------
 python/teams/team-one/src/team_one/messages.py |  1 +
 3 files changed, 9 insertions(+), 15 deletions(-)

diff --git a/python/benchmarks/HumanEval/Templates/TeamOne/scenario.py b/python/benchmarks/HumanEval/Templates/TeamOne/scenario.py
index 4dfb74365..d37412d19 100644
--- a/python/benchmarks/HumanEval/Templates/TeamOne/scenario.py
+++ b/python/benchmarks/HumanEval/Templates/TeamOne/scenario.py
@@ -89,7 +89,10 @@ class MyHandler(logging.Handler):
     def emit(self, record: logging.LogRecord) -> None:
         try:
             if isinstance(record.msg, OrchestrationEvent):
-                print(record.msg.message, flush=True)
+                print(f"""---------------------------------------------------------------------------
+\033[91m{record.msg.source}:\033[0m
+
+{record.msg.message}""", flush=True)
         except Exception:
             self.handleError(record)
 
diff --git a/python/teams/team-one/src/team_one/agents/orchestrator.py b/python/teams/team-one/src/team_one/agents/orchestrator.py
index e830f2eb7..93fb442ea 100644
--- a/python/teams/team-one/src/team_one/agents/orchestrator.py
+++ b/python/teams/team-one/src/team_one/agents/orchestrator.py
@@ -29,18 +29,10 @@ class RoundRobinOrchestrator(TypeRoutedAgent):
         if isinstance(message.content, UserMessage) or isinstance(message.content, AssistantMessage):
             source = message.content.source
 
-        assert isinstance(source, str)
+        content = str(message.content.content)
 
         current_timestamp = datetime.now().isoformat()
-        logger.info(
-            OrchestrationEvent(
-                current_timestamp,
-                f"""{source}:
-
-{message.content.content}
--------------------------------------""",
-            )
-        )
+        logger.info(OrchestrationEvent(current_timestamp, source, content))
 
         if self._num_rounds > 20:
             return
@@ -53,10 +45,8 @@ class RoundRobinOrchestrator(TypeRoutedAgent):
         logger.info(
             OrchestrationEvent(
                 current_timestamp,
-                f"""Orchestrator (thought):
-
-Next speaker {next_agent.metadata['name']}
--------------------------------------""",
+                source="Orchestrator (thought)",
+                message=f"Next speaker {next_agent.metadata['name']}" "",
             )
         )
 
diff --git a/python/teams/team-one/src/team_one/messages.py b/python/teams/team-one/src/team_one/messages.py
index df207ac98..9a0e738cc 100644
--- a/python/teams/team-one/src/team_one/messages.py
+++ b/python/teams/team-one/src/team_one/messages.py
@@ -16,4 +16,5 @@ class RequestReplyMessage:
 @dataclass
 class OrchestrationEvent:
     timestamp: str
+    source: str
     message: str