mirror of
https://github.com/All-Hands-AI/OpenHands.git
synced 2026-01-10 15:28:14 -05:00
CLI: custom visualizer (#11677)
This commit is contained in:
Rohit Malhotra
312
openhands-cli/openhands_cli/tui/visualizer.py
Normal file
312
openhands-cli/openhands_cli/tui/visualizer.py
Normal file
@@ -0,0 +1,312 @@
|
||||
import re
|
||||
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
from openhands.sdk.conversation.visualizer.base import (
|
||||
ConversationVisualizerBase,
|
||||
)
|
||||
from openhands.sdk.event import (
|
||||
ActionEvent,
|
||||
AgentErrorEvent,
|
||||
MessageEvent,
|
||||
ObservationEvent,
|
||||
PauseEvent,
|
||||
SystemPromptEvent,
|
||||
UserRejectObservation,
|
||||
)
|
||||
from openhands.sdk.event.base import Event
|
||||
from openhands.sdk.event.condenser import Condensation
|
||||
|
||||
|
||||
# These are external inputs
|
||||
_OBSERVATION_COLOR = "yellow"
|
||||
_MESSAGE_USER_COLOR = "gold3"
|
||||
_PAUSE_COLOR = "bright_yellow"
|
||||
# These are internal system stuff
|
||||
_SYSTEM_COLOR = "magenta"
|
||||
_THOUGHT_COLOR = "bright_black"
|
||||
_ERROR_COLOR = "red"
|
||||
# These are agent actions
|
||||
_ACTION_COLOR = "blue"
|
||||
_MESSAGE_ASSISTANT_COLOR = _ACTION_COLOR
|
||||
|
||||
DEFAULT_HIGHLIGHT_REGEX = {
|
||||
r"^Reasoning:": f"bold {_THOUGHT_COLOR}",
|
||||
r"^Thought:": f"bold {_THOUGHT_COLOR}",
|
||||
r"^Action:": f"bold {_ACTION_COLOR}",
|
||||
r"^Arguments:": f"bold {_ACTION_COLOR}",
|
||||
r"^Tool:": f"bold {_OBSERVATION_COLOR}",
|
||||
r"^Result:": f"bold {_OBSERVATION_COLOR}",
|
||||
r"^Rejection Reason:": f"bold {_ERROR_COLOR}",
|
||||
# Markdown-style
|
||||
r"\*\*(.*?)\*\*": "bold",
|
||||
r"\*(.*?)\*": "italic",
|
||||
}
|
||||
|
||||
_PANEL_PADDING = (1, 1)
|
||||
|
||||
|
||||
class CLIVisualizer(ConversationVisualizerBase):
|
||||
"""Handles visualization of conversation events with Rich formatting.
|
||||
|
||||
Provides Rich-formatted output with panels and complete content display.
|
||||
"""
|
||||
|
||||
_console: Console
|
||||
_skip_user_messages: bool
|
||||
_highlight_patterns: dict[str, str]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str | None = None,
|
||||
highlight_regex: dict[str, str] | None = DEFAULT_HIGHLIGHT_REGEX,
|
||||
skip_user_messages: bool = False,
|
||||
):
|
||||
"""Initialize the visualizer.
|
||||
|
||||
Args:
|
||||
name: Optional name to prefix in panel titles to identify
|
||||
which agent/conversation is speaking.
|
||||
highlight_regex: Dictionary mapping regex patterns to Rich color styles
|
||||
for highlighting keywords in the visualizer.
|
||||
For example: {"Reasoning:": "bold blue",
|
||||
"Thought:": "bold green"}
|
||||
skip_user_messages: If True, skip displaying user messages. Useful for
|
||||
scenarios where user input is not relevant to show.
|
||||
"""
|
||||
super().__init__(
|
||||
name=name,
|
||||
)
|
||||
self._console = Console()
|
||||
self._skip_user_messages = skip_user_messages
|
||||
self._highlight_patterns = highlight_regex or {}
|
||||
|
||||
def on_event(self, event: Event) -> None:
|
||||
"""Main event handler that displays events with Rich formatting."""
|
||||
panel = self._create_event_panel(event)
|
||||
if panel:
|
||||
self._console.print(panel)
|
||||
self._console.print() # Add spacing between events
|
||||
|
||||
def _apply_highlighting(self, text: Text) -> Text:
|
||||
"""Apply regex-based highlighting to text content.
|
||||
|
||||
Args:
|
||||
text: The Rich Text object to highlight
|
||||
|
||||
Returns:
|
||||
A new Text object with highlighting applied
|
||||
"""
|
||||
if not self._highlight_patterns:
|
||||
return text
|
||||
|
||||
# Create a copy to avoid modifying the original
|
||||
highlighted = text.copy()
|
||||
|
||||
# Apply each pattern using Rich's built-in highlight_regex method
|
||||
for pattern, style in self._highlight_patterns.items():
|
||||
pattern_compiled = re.compile(pattern, re.MULTILINE)
|
||||
highlighted.highlight_regex(pattern_compiled, style)
|
||||
|
||||
return highlighted
|
||||
|
||||
def _create_event_panel(self, event: Event) -> Panel | None:
|
||||
"""Create a Rich Panel for the event with appropriate styling."""
|
||||
# Use the event's visualize property for content
|
||||
content = event.visualize
|
||||
|
||||
if not content.plain.strip():
|
||||
return None
|
||||
|
||||
# Apply highlighting if configured
|
||||
if self._highlight_patterns:
|
||||
content = self._apply_highlighting(content)
|
||||
|
||||
# Don't emit system prompt in CLI
|
||||
if isinstance(event, SystemPromptEvent):
|
||||
title = f"[bold {_SYSTEM_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
title += f"System Prompt[/bold {_SYSTEM_COLOR}]"
|
||||
return None
|
||||
elif isinstance(event, ActionEvent):
|
||||
# Check if action is None (non-executable)
|
||||
title = f"[bold {_ACTION_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
if event.action is None:
|
||||
title += f"Agent Action (Not Executed)[/bold {_ACTION_COLOR}]"
|
||||
else:
|
||||
title += f"Agent Action[/bold {_ACTION_COLOR}]"
|
||||
return Panel(
|
||||
content,
|
||||
title=title,
|
||||
subtitle=self._format_metrics_subtitle(),
|
||||
border_style=_ACTION_COLOR,
|
||||
padding=_PANEL_PADDING,
|
||||
expand=True,
|
||||
)
|
||||
elif isinstance(event, ObservationEvent):
|
||||
title = f"[bold {_OBSERVATION_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
title += f"Observation[/bold {_OBSERVATION_COLOR}]"
|
||||
return Panel(
|
||||
content,
|
||||
title=title,
|
||||
border_style=_OBSERVATION_COLOR,
|
||||
padding=_PANEL_PADDING,
|
||||
expand=True,
|
||||
)
|
||||
elif isinstance(event, UserRejectObservation):
|
||||
title = f"[bold {_ERROR_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
title += f"User Rejected Action[/bold {_ERROR_COLOR}]"
|
||||
return Panel(
|
||||
content,
|
||||
title=title,
|
||||
border_style=_ERROR_COLOR,
|
||||
padding=_PANEL_PADDING,
|
||||
expand=True,
|
||||
)
|
||||
elif isinstance(event, MessageEvent):
|
||||
if (
|
||||
self._skip_user_messages
|
||||
and event.llm_message
|
||||
and event.llm_message.role == "user"
|
||||
):
|
||||
return
|
||||
assert event.llm_message is not None
|
||||
# Role-based styling
|
||||
role_colors = {
|
||||
"user": _MESSAGE_USER_COLOR,
|
||||
"assistant": _MESSAGE_ASSISTANT_COLOR,
|
||||
}
|
||||
role_color = role_colors.get(event.llm_message.role, "white")
|
||||
|
||||
# "User Message To [Name] Agent" for user
|
||||
# "Message from [Name] Agent" for agent
|
||||
agent_name = f"{self._name} " if self._name else ""
|
||||
|
||||
if event.llm_message.role == "user":
|
||||
title_text = (
|
||||
f"[bold {role_color}]User Message to "
|
||||
f"{agent_name}Agent[/bold {role_color}]"
|
||||
)
|
||||
else:
|
||||
title_text = (
|
||||
f"[bold {role_color}]Message from "
|
||||
f"{agent_name}Agent[/bold {role_color}]"
|
||||
)
|
||||
return Panel(
|
||||
content,
|
||||
title=title_text,
|
||||
subtitle=self._format_metrics_subtitle(),
|
||||
border_style=role_color,
|
||||
padding=_PANEL_PADDING,
|
||||
expand=True,
|
||||
)
|
||||
elif isinstance(event, AgentErrorEvent):
|
||||
title = f"[bold {_ERROR_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
title += f"Agent Error[/bold {_ERROR_COLOR}]"
|
||||
return Panel(
|
||||
content,
|
||||
title=title,
|
||||
subtitle=self._format_metrics_subtitle(),
|
||||
border_style=_ERROR_COLOR,
|
||||
padding=_PANEL_PADDING,
|
||||
expand=True,
|
||||
)
|
||||
elif isinstance(event, PauseEvent):
|
||||
title = f"[bold {_PAUSE_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
title += f"User Paused[/bold {_PAUSE_COLOR}]"
|
||||
return Panel(
|
||||
content,
|
||||
title=title,
|
||||
border_style=_PAUSE_COLOR,
|
||||
padding=_PANEL_PADDING,
|
||||
expand=True,
|
||||
)
|
||||
elif isinstance(event, Condensation):
|
||||
title = f"[bold {_SYSTEM_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
title += f"Condensation[/bold {_SYSTEM_COLOR}]"
|
||||
return Panel(
|
||||
content,
|
||||
title=title,
|
||||
subtitle=self._format_metrics_subtitle(),
|
||||
border_style=_SYSTEM_COLOR,
|
||||
expand=True,
|
||||
)
|
||||
else:
|
||||
# Fallback panel for unknown event types
|
||||
title = f"[bold {_ERROR_COLOR}]"
|
||||
if self._name:
|
||||
title += f"{self._name} "
|
||||
title += f"UNKNOWN Event: {event.__class__.__name__}[/bold {_ERROR_COLOR}]"
|
||||
return Panel(
|
||||
content,
|
||||
title=title,
|
||||
subtitle=f"({event.source})",
|
||||
border_style=_ERROR_COLOR,
|
||||
padding=_PANEL_PADDING,
|
||||
expand=True,
|
||||
)
|
||||
|
||||
def _format_metrics_subtitle(self) -> str | None:
|
||||
"""Format LLM metrics as a visually appealing subtitle string with icons,
|
||||
colors, and k/m abbreviations using conversation stats."""
|
||||
stats = self.conversation_stats
|
||||
if not stats:
|
||||
return None
|
||||
|
||||
combined_metrics = stats.get_combined_metrics()
|
||||
if not combined_metrics or not combined_metrics.accumulated_token_usage:
|
||||
return None
|
||||
|
||||
usage = combined_metrics.accumulated_token_usage
|
||||
cost = combined_metrics.accumulated_cost or 0.0
|
||||
|
||||
# helper: 1234 -> "1.2K", 1200000 -> "1.2M"
|
||||
def abbr(n: int | float) -> str:
|
||||
n = int(n or 0)
|
||||
if n >= 1_000_000_000:
|
||||
val, suffix = n / 1_000_000_000, "B"
|
||||
elif n >= 1_000_000:
|
||||
val, suffix = n / 1_000_000, "M"
|
||||
elif n >= 1_000:
|
||||
val, suffix = n / 1_000, "K"
|
||||
else:
|
||||
return str(n)
|
||||
return f"{val:.2f}".rstrip("0").rstrip(".") + suffix
|
||||
|
||||
input_tokens = abbr(usage.prompt_tokens or 0)
|
||||
output_tokens = abbr(usage.completion_tokens or 0)
|
||||
|
||||
# Cache hit rate (prompt + cache)
|
||||
prompt = usage.prompt_tokens or 0
|
||||
cache_read = usage.cache_read_tokens or 0
|
||||
cache_rate = f"{(cache_read / prompt * 100):.2f}%" if prompt > 0 else "N/A"
|
||||
reasoning_tokens = usage.reasoning_tokens or 0
|
||||
|
||||
# Cost
|
||||
cost_str = f"{cost:.4f}" if cost > 0 else "0.00"
|
||||
|
||||
# Build with fixed color scheme
|
||||
parts: list[str] = []
|
||||
parts.append(f"[cyan]↑ input {input_tokens}[/cyan]")
|
||||
parts.append(f"[magenta]cache hit {cache_rate}[/magenta]")
|
||||
if reasoning_tokens > 0:
|
||||
parts.append(f"[yellow] reasoning {abbr(reasoning_tokens)}[/yellow]")
|
||||
parts.append(f"[blue]↓ output {output_tokens}[/blue]")
|
||||
parts.append(f"[green]$ {cost_str}[/green]")
|
||||
|
||||
return "Tokens: " + " • ".join(parts)
|
||||
Reference in New Issue
Block a user