File size: 2,651 Bytes
363cda9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import json
from rich.console import Console
from rich.syntax import Syntax
from rich.panel import Panel
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
console = Console(width=100, soft_wrap=False)
_last_context_snapshot = None # used to suppress repeated context
_last_message_ids = set() # track printed messages
def serialize_message(msg) -> dict:
"""Convert LangChain message objects into serializable dicts."""
if hasattr(msg, "dict"):
return msg.dict()
elif hasattr(msg, "__dict__"):
return {k: serialize_message(v) for k, v in msg.__dict__.items()}
elif isinstance(msg, list):
return [serialize_message(v) for v in msg]
elif isinstance(msg, dict):
return {k: serialize_message(v) for k, v in msg.items()}
else:
return msg
def pretty_print_state(state: dict, show_context: bool = True) -> None:
"""
Pretty-print the agent's state in a clean, color-coded way.
Parameters
----------
state : dict
The LangGraph agent state chunk (from the stream).
show_context : bool, optional
Whether to display the context (default True).
If True, only shows context when it has changed since last call.
"""
global _last_context_snapshot
# --- Display message chunks ---
for msg in state.get("messages", []):
msg_id = getattr(msg, "id", id(msg))
if msg_id in _last_message_ids:
continue # skip duplicates
_last_message_ids.add(msg_id)
msg_dict = serialize_message(msg)
msg_json = json.dumps(msg_dict, indent=2)
if isinstance(msg, HumanMessage):
color, title = "cyan", "🧑 HumanMessage"
elif isinstance(msg, ToolMessage):
color, title = "yellow", f"🧰 ToolMessage ({msg_dict.get('name','?')})"
elif isinstance(msg, AIMessage):
color, title = "magenta", "🤖 AIMessage"
else:
color, title = "white", "Other"
syntax = Syntax(msg_json, "json", theme="monokai", line_numbers=False)
console.print(Panel(syntax, title=title, border_style=color))
# --- Optional context view ---
#if show_context:
# context = state.get("context", {})
# if context and context != _last_context_snapshot:
# _last_context_snapshot = context.copy() # cache for next comparison
# context_json = json.dumps(context, indent=2, default=str)
# syntax = Syntax(context_json, "json", theme="monokai", line_numbers=False)
# console.print(Panel(syntax, title="🧠 Context (updated)", border_style="green"))
|