Skip to content

Commit

Permalink
Feature/seperate render for humans (#24)
Browse files Browse the repository at this point in the history 
* seperate the functions of render for humans

* refactor codes
  • Loading branch information
@XuhuiZhou
XuhuiZhou authored Jul 28, 2024
1 parent 0575361 commit 801734c
Show file tree
Hide file tree
Showing 4 changed files with 155 additions and 109 deletions.
2 changes: 2 additions & 0 deletions haicosystem/protocols/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
HaiScriptBackground,
LangchainAgentAction,
SimulatedObservation,
messageForRendering,
)

__all__ = [
Expand All @@ -12,4 +13,5 @@
"LangchainAgentAction",
"HaiScriptBackground",
"HaiAgentAction",
"messageForRendering",
]
12 changes: 8 additions & 4 deletions haicosystem/protocols/messages.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
from typing import TypedDict

from langchain.schema import AgentAction as LAgentAction
from langchain_core.load.serializable import (
SerializedConstructor,
SerializedNotImplemented,
)
from pydantic import Field
from sotopia.messages import ActionType, AgentAction, Message, ScriptBackground
from sotopia.utils import format_docstring
Expand Down Expand Up @@ -72,3 +70,9 @@ class HaiAgentAction(AgentAction):
argument: str = Field(
description="the utterance if choose 'speak' (which should be a string instead of a JSON), the expression or gesture if choose 'non-verbal communication', or the tool calling action if choose 'action'" # TODO: assumption whenerver the action_type is 'action', the argument is the tool calling action
)


class messageForRendering(TypedDict):
role: str
type: str
content: str
244 changes: 141 additions & 103 deletions haicosystem/utils/render.py
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
from rich import print
from typing import TypedDict

from rich.console import Console
from rich.json import JSON
from rich.panel import Panel
from sotopia.database import EpisodeLog

from haicosystem.protocols import messageForRendering


def pick_color_for_agent(
agent_name: str,
Expand Down Expand Up @@ -45,17 +49,11 @@ def parse_reasoning(reasoning: str, num_agents: int) -> tuple[list[str], str]:
return comment_chunks, general_comment


def render_for_humans(episode: EpisodeLog) -> None:
"""Generate a human readable version of the episode log."""
def render_for_humans(episode: EpisodeLog) -> list[messageForRendering]:
"""Generate a list of messages for human-readable version of the episode log."""

messages_for_rendering: list[messageForRendering] = []

available_colors: list[str] = [
"medium_violet_red",
"green",
"slate_blue1",
"yellow",
]
occupied_colors: list[str] = []
agent_color_map: dict[str, str] = {}
for idx, turn in enumerate(episode.messages):
is_observation_printed = False

Expand All @@ -64,38 +62,26 @@ def render_for_humans(episode: EpisodeLog) -> None:
len(turn) >= 2
), "The first turn should have at least environment messages"

print(
Panel(
turn[0][2],
title="Background Info",
style="blue",
title_align="left",
)
messages_for_rendering.append(
{"role": "Background Info", "type": "info", "content": turn[0][2]}
)
print(
Panel(
turn[1][2],
title="Background Info",
style="blue",
title_align="left",
)
messages_for_rendering.append(
{"role": "Background Info", "type": "info", "content": turn[1][2]}
)
messages_for_rendering.append(
{"role": "System", "type": "divider", "content": "Start Simulation"}
)
print("=" * 100)
print("Start Simulation")
print("=" * 100)

for sender, receiver, message in turn:
# "Observation" indicates the agent takes an action
if not is_observation_printed and "Observation:" in message and idx != 0:
extract_observation = message.split("Observation:")[1].strip()
if extract_observation:
print(
Panel(
JSON(extract_observation, highlight=False),
title="Observation",
style="yellow",
title_align="left",
)
messages_for_rendering.append(
{
"role": "Observation",
"type": "observation",
"content": extract_observation,
}
)
is_observation_printed = True

Expand All @@ -104,87 +90,139 @@ def render_for_humans(episode: EpisodeLog) -> None:
if "did nothing" in message:
continue
elif "left the conversation" in message:
print(
Panel(
f"{sender} left the conversation",
title="Environment",
style="red",
title_align="left",
)
messages_for_rendering.append(
{
"role": "Environment",
"type": "leave",
"content": f"{sender} left the conversation",
}
)
else:
# Conversation
if "said:" in message:
(
available_colors,
occupied_colors,
agent_color_map,
sender_color,
) = pick_color_for_agent(
sender,
available_colors,
occupied_colors,
agent_color_map,
)
message = message.split("said:")[1].strip()
print(
Panel(
f"{message}",
title=f"{sender} (Said)",
style=sender_color,
title_align="left",
)
messages_for_rendering.append(
{"role": sender, "type": "said", "content": message}
)
# Action
else:
(
available_colors,
occupied_colors,
agent_color_map,
sender_color,
) = pick_color_for_agent(
sender,
available_colors,
occupied_colors,
agent_color_map,
)
message = message.replace("[action]", "")
print(
Panel(
JSON(message, highlight=False),
title=f"{sender} (Action)",
style=sender_color,
title_align="left",
)
messages_for_rendering.append(
{"role": sender, "type": "action", "content": message}
)
else:
print(Panel(message, style="white"))
print("=" * 100)
print("End Simulation")
print("=" * 100)
messages_for_rendering.append(
{
"role": "Environment",
"type": "environment",
"content": message,
}
)

messages_for_rendering.append(
{"role": "System", "type": "divider", "content": "End Simulation"}
)

reasoning_per_agent, general_comment = parse_reasoning(
episode.reasoning, len(agent_color_map)
episode.reasoning,
len(
set(
msg["role"]
for msg in messages_for_rendering
if msg["type"] in {"said", "action"}
)
),
)

print(
Panel(
general_comment,
title="General (Comments)",
style="blue",
title_align="left",
)
messages_for_rendering.append(
{"role": "General", "type": "comment", "content": general_comment}
)

for idx, reasoning in enumerate(reasoning_per_agent):
print(
Panel(
f"{reasoning}\n"
+ "=" * 100
+ "\nRewards: "
+ str(episode.rewards[idx]),
title=f"Agent {idx + 1} (Comments)",
style="blue",
title_align="left",
)
messages_for_rendering.append(
{
"role": f"Agent {idx + 1}",
"type": "comment",
"content": f"{reasoning}\n{'=' * 100}\nRewards: {str(episode.rewards[idx])}",
}
)

return messages_for_rendering


def rich_rendering(messages: list[messageForRendering]) -> None:
"""Render the list of messages using rich library."""

console = Console()

available_colors: list[str] = [
"medium_violet_red",
"green",
"slate_blue1",
"yellow",
]
occupied_colors: list[str] = []
agent_color_map: dict[str, str] = {}

def pick_color_for_agent(agent: str) -> str:
if agent not in agent_color_map:
if available_colors:
color = available_colors.pop(0)
occupied_colors.append(color)
else:
color = occupied_colors.pop(0)
available_colors.append(color)
agent_color_map[agent] = color
return agent_color_map[agent]

for message in messages:
role = message["role"]
msg_type = message["type"]
content = message["content"]

if msg_type == "info":
console.print(Panel(content, title=role, style="blue", title_align="left"))
elif msg_type == "divider":
console.print("=" * 100)
console.print(content)
console.print("=" * 100)
elif msg_type == "observation":
console.print(
Panel(
JSON(content, highlight=False),
title=role,
style="yellow",
title_align="left",
)
)
elif msg_type == "leave":
console.print(Panel(content, title=role, style="red", title_align="left"))
elif msg_type == "said":
sender_color = pick_color_for_agent(role)
console.print(
Panel(
content,
title=f"{role} (Said)",
style=sender_color,
title_align="left",
)
)
elif msg_type == "action":
sender_color = pick_color_for_agent(role)
console.print(
Panel(
JSON(content, highlight=False),
title=f"{role} (Action)",
style=sender_color,
title_align="left",
)
)
elif msg_type == "environment":
console.print(Panel(content, style="white"))
elif msg_type == "comment":
console.print(
Panel(
content,
title=f"{role} (Comments)",
style="blue",
title_align="left",
)
)
6 changes: 4 additions & 2 deletions notebooks/render_for_human.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from sotopia.database import EpisodeLog
from typer import Typer

from haicosystem.utils.render import render_for_humans
from haicosystem.utils.render import render_for_humans, rich_rendering

app = Typer()

Expand All @@ -13,7 +13,9 @@ def print_episode(episode_number: int) -> None:
"""
episode = EpisodeLog.find(EpisodeLog.tag == "haicosystem_debug")[episode_number] # type: ignore
assert isinstance(episode, EpisodeLog)
render_for_humans(episode)
messages = render_for_humans(episode)
breakpoint()
rich_rendering(messages)


if __name__ == "__main__":
Expand Down

0 comments on commit 801734c

Please sign in to comment.