feat: add function to print communication between agents

utils/__init__.py

@@ -0,0 +1,3 @@
+from utils.pretty_print_messages import pretty_print_messages
+
+__all__ = ["pretty_print_messages"]

utils/pretty_print_messages.py

@@ -0,0 +1,46 @@
+from typing import Any, Dict, List, Tuple, Union
+
+from langchain_core.messages import BaseMessage, convert_to_messages
+
+
+def pretty_print_message(message: BaseMessage, indent: bool = False) -> None:
+    pretty_message = message.pretty_repr(html=True)
+    if not indent:
+        print(pretty_message)
+        return
+
+    indented = "\n".join("\t" + c for c in pretty_message.split("\n"))
+    print(indented)
+
+
+def pretty_print_messages(
+    update: Union[Dict[str, Any], Tuple[List[str], Dict[str, Any]]],
+    last_message: bool = False,
+) -> None:
+    is_subgraph = False
+    if isinstance(update, tuple):
+        ns, update = update
+        # skip parent graph updates in the printouts
+        if len(ns) == 0:
+            return
+
+        graph_id = ns[-1].split(":")[0]
+        print(f"Update from subgraph {graph_id}:")
+        print("\n")
+        is_subgraph = True
+
+    for node_name, node_update in update.items():
+        update_label = f"Update from node {node_name}:"
+        if is_subgraph:
+            update_label = "\t" + update_label
+
+        print(update_label)
+        print("\n")
+
+        messages = convert_to_messages(node_update["messages"])
+        if last_message:
+            messages = messages[-1:]
+
+        for m in messages:
+            pretty_print_message(m, indent=is_subgraph)
+        print("\n")