diff --git a/src/sherpa_ai/__init__.py b/src/sherpa_ai/__init__.py index 3dc1f76bc..8bb4dd31e 100644 --- a/src/sherpa_ai/__init__.py +++ b/src/sherpa_ai/__init__.py @@ -1 +1,9 @@ -__version__ = "0.1.0" +"""Visualization tools for Sherpa AI. + +This package provides visualization utilities for inspecting agent behavior, +state machine execution, and decision trajectories. +""" + +from sherpa_ai.visualization.state_machine_viewer import StateMachineViewer + +__all__ = ["StateMachineViewer"] \ No newline at end of file diff --git a/src/sherpa_ai/agents/base.py b/src/sherpa_ai/agents/base.py index 5f7017dfc..33176f2a1 100644 --- a/src/sherpa_ai/agents/base.py +++ b/src/sherpa_ai/agents/base.py @@ -400,6 +400,34 @@ async def async_run(self) -> TaskResult: result = await self.async_select_action() + # Record the decision: what was chosen and what was skipped + if result is not None and not isinstance(result, Exception): + chosen_name = result.action.name + alternative_names = [ + a.name for a in actions if a.name != chosen_name + ] + current_state = ( + self.belief.get_state() + if self.belief.state_machine + else "" + ) + self.belief.update_internal( + "decision", + self.name, + chosen=chosen_name, + alternatives=alternative_names, + state=current_state or "", + ) + if self.shared_memory is not None: + await self.shared_memory.async_add( + "decision", + self.name, + sender=self.name, + chosen=chosen_name, + alternatives=alternative_names, + state=current_state or "", + ) + if result is None: # this means no action is selected continue diff --git a/src/sherpa_ai/events.py b/src/sherpa_ai/events.py index a0e8db809..3a0953faf 100644 --- a/src/sherpa_ai/events.py +++ b/src/sherpa_ai/events.py @@ -173,6 +173,30 @@ class ActionFinishEvent(Event): event_type: str = Field("action_finish", frozen=True) +class DecisionEvent(Event): + """Event recording an action selection decision. + + Captures which action was chosen from the available candidates, + along with the alternatives that were not taken. Used for + trajectory analysis and state machine visualization. + + Attributes: + name (str): Name of the decision point. + chosen (str): Name of the action that was selected. + alternatives (list[str]): Names of actions that were available + but not selected. + state (str): The state machine state at decision time, or empty + string if no state machine is used. + event_type (str): Fixed to "decision". + """ + + name: str + chosen: str + alternatives: list[str] + state: str = "" + event_type: str = Field("decision", frozen=True) + + def build_event(event_type: str, name: str, **kwargs) -> Event: """Factory function to create appropriate Event objects based on event type. @@ -193,6 +217,7 @@ def build_event(event_type: str, name: str, **kwargs) -> Event: Event: An instance of the appropriate Event subclass: - ActionStartEvent for "action_start" event_type - ActionFinishEvent for "action_finish" event_type + - DecisionEvent for "decision" event_type - GenericEvent for all other event_type values Example: @@ -206,10 +231,12 @@ def build_event(event_type: str, name: str, **kwargs) -> Event: return ActionStartEvent(**kwargs, name=name) elif event_type == "action_finish": return ActionFinishEvent(**kwargs, name=name) + elif event_type == "decision": + return DecisionEvent(name=name, **kwargs) elif event_type == "trigger": if not kwargs.get("args"): # Default to no args if not provided kwargs["args"] = {} return TriggerEvent(name=name, **kwargs) else: - return GenericEvent(**kwargs, event_type=event_type, name=name) + return GenericEvent(**kwargs, event_type=event_type, name=name) \ No newline at end of file diff --git a/src/sherpa_ai/visualization/__init__.py b/src/sherpa_ai/visualization/__init__.py new file mode 100644 index 000000000..8bb4dd31e --- /dev/null +++ b/src/sherpa_ai/visualization/__init__.py @@ -0,0 +1,9 @@ +"""Visualization tools for Sherpa AI. + +This package provides visualization utilities for inspecting agent behavior, +state machine execution, and decision trajectories. +""" + +from sherpa_ai.visualization.state_machine_viewer import StateMachineViewer + +__all__ = ["StateMachineViewer"] \ No newline at end of file diff --git a/src/sherpa_ai/visualization/state_machine_viewer.py b/src/sherpa_ai/visualization/state_machine_viewer.py new file mode 100644 index 000000000..1b642466c --- /dev/null +++ b/src/sherpa_ai/visualization/state_machine_viewer.py @@ -0,0 +1,306 @@ + +"""State machine visualization module for Sherpa AI. + +Reads DecisionEvents from an agent's belief and renders the state machine +as a standalone HTML file showing states, taken transitions, and untaken +alternatives at each decision step. + +Usage: + >>> from sherpa_ai.visualization import StateMachineViewer + >>> viewer = StateMachineViewer(belief=agent.belief) + >>> viewer.render("trajectory.html") +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import TYPE_CHECKING, Optional + +from loguru import logger + +if TYPE_CHECKING: + from sherpa_ai.memory.belief import Belief + from sherpa_ai.memory.state_machine import SherpaStateMachine + + +class DecisionRecord: + """A single decision point extracted from a DecisionEvent.""" + + def __init__(self, step: int, state: str, chosen: str, alternatives: list[str]): + self.step = step + self.state = state + self.chosen = chosen + self.alternatives = alternatives + + +class StateMachineViewer: + """Generates an HTML visualization of state machine execution.""" + + def __init__(self, belief: Belief): + self.belief = belief + self.state_machine = belief.state_machine + + def extract_decisions(self) -> list[DecisionRecord]: + """Extract decision records from the belief's internal events.""" + decisions = [] + step = 0 + for event in self.belief.internal_events: + if event.event_type == "decision": + step += 1 + decisions.append( + DecisionRecord( + step=step, + state=event.state, + chosen=event.chosen, + alternatives=event.alternatives, + ) + ) + return decisions + + def extract_graph(self) -> dict: + """Extract the state machine graph structure.""" + if self.state_machine is None: + logger.warning("No state machine attached to belief") + return {"states": [], "transitions": []} + + sm = self.state_machine.sm + states = [s.name for s in sm.states.values()] + + transitions = [] + for trigger_name, event in sm.events.items(): + if trigger_name not in self.state_machine.explicit_transitions: + continue + for source, trans_list in event.transitions.items(): + for t in trans_list: + transitions.append( + { + "trigger": trigger_name, + "source": source, + "dest": t.dest if t.dest else source, + } + ) + + return {"states": states, "transitions": transitions} + + def render(self, output_path: str = "state_machine.html") -> str: + """Generate an HTML visualization and write it to a file.""" + graph = self.extract_graph() + decisions = self.extract_decisions() + current_state = self.belief.get_state() if self.state_machine else None + + html = self._build_html(graph, decisions, current_state) + + path = Path(output_path) + path.write_text(html, encoding="utf-8") + logger.info(f"State machine visualization written to {path.absolute()}") + return str(path.absolute()) + + def render_html(self) -> str: + """Generate the HTML visualization as a string without writing to file.""" + graph = self.extract_graph() + decisions = self.extract_decisions() + current_state = self.belief.get_state() if self.state_machine else None + return self._build_html(graph, decisions, current_state) + + def _build_html( + self, + graph: dict, + decisions: list[DecisionRecord], + current_state: Optional[str], + ) -> str: + states = graph["states"] + transitions = graph["transitions"] + + taken_triggers = {d.chosen for d in decisions} + skipped_triggers = set() + for d in decisions: + skipped_triggers.update(d.alternatives) + skipped_triggers -= taken_triggers + + node_positions = self._layout_circle(states, cx=400, cy=300, radius=200) + + svg_arrows = self._build_arrows( + transitions, node_positions, taken_triggers, skipped_triggers + ) + svg_nodes = self._build_nodes(states, node_positions, current_state) + decision_table = self._build_decision_table(decisions) + + html = f""" + + + +Sherpa State Machine Viewer + + + +
+

Sherpa State Machine Viewer

+
+ + + + + + + + + + + + + {svg_arrows} + {svg_nodes} + +
+
Taken
+
Skipped
+
Not yet reached
+
Current state
+
+
+

Decision History

+ {decision_table} +
+ +""" + return html + + def _layout_circle(self, states, cx, cy, radius): + positions = {} + n = len(states) + if n == 0: + return positions + for i, state in enumerate(states): + angle = (2 * math.pi * i / n) - math.pi / 2 + x = cx + radius * math.cos(angle) + y = cy + radius * math.sin(angle) + positions[state] = (x, y) + return positions + + def _build_nodes(self, states, positions, current_state): + elements = [] + for state in states: + x, y = positions.get(state, (400, 300)) + is_current = state == current_state + fill = "#238636" if is_current else "#1f6feb" + stroke = "#56d364" if is_current else "#388bfd" + r = 36 if is_current else 32 + elements.append( + f'' + ) + elements.append( + f'{state}' + ) + return "\n ".join(elements) + + def _build_arrows(self, transitions, positions, taken, skipped): + elements = [] + edge_counts: dict[tuple[str, str], int] = {} + for t in transitions: + source, dest, trigger = t["source"], t["dest"], t["trigger"] + if source not in positions or dest not in positions: + continue + sx, sy = positions[source] + dx, dy = positions[dest] + if trigger in taken: + color, dash, width, marker = "#3fb950", "", "2.5", "url(#arrow-taken)" + elif trigger in skipped: + color, dash, width, marker = "#484f58", 'stroke-dasharray="6,4"', "1.5", "url(#arrow-skipped)" + else: + color, dash, width, marker = "#58a6ff", 'stroke-dasharray="2,3"', "1", "url(#arrow-default)" + edge_key = (min(source, dest), max(source, dest)) + count = edge_counts.get(edge_key, 0) + edge_counts[edge_key] = count + 1 + if source == dest: + elements.append( + f'' + ) + else: + node_r = 34 + dx_vec, dy_vec = dx - sx, dy - sy + length = math.sqrt(dx_vec**2 + dy_vec**2) + if length == 0: + continue + ux, uy = dx_vec / length, dy_vec / length + start_x, start_y = sx + ux * node_r, sy + uy * node_r + end_x, end_y = dx - ux * node_r, dy - uy * node_r + offset = (count - 0.5) * 20 if count > 0 else 0 + nx, ny = -uy * offset, ux * offset + if offset == 0: + elements.append( + f'' + ) + else: + mid_x = (start_x + end_x) / 2 + nx + mid_y = (start_y + end_y) / 2 + ny + elements.append( + f'' + ) + label_x = (start_x + end_x) / 2 + nx * 0.6 + label_y = (start_y + end_y) / 2 + ny * 0.6 - 6 + elements.append( + f'{trigger}' + ) + return "\n ".join(elements) + + def _build_decision_table(self, decisions): + if not decisions: + return '
No decisions recorded yet.
' + rows = [] + for d in decisions: + alts = ", ".join(d.alternatives) if d.alternatives else "none" + rows.append( + f"{d.step}" + f'{d.state or "—"}' + f'{d.chosen}' + f'{alts}' + ) + return f""" + + {"".join(rows)} +
StepStateChosenAlternatives
""" \ No newline at end of file diff --git a/src/tests/unit_tests/policies/test_chat_based_policy.py b/src/tests/unit_tests/policies/test_chat_based_policy.py index 09c345d72..95e248342 100644 --- a/src/tests/unit_tests/policies/test_chat_based_policy.py +++ b/src/tests/unit_tests/policies/test_chat_based_policy.py @@ -98,5 +98,5 @@ def test_chat_based_policy_with_agent(get_llm): # noqa: F811 # The agent should be able to move to state D assert result.status == "success" - assert len(belief.internal_events) == 4 + assert len(belief.internal_events) == 6 assert belief.get_state() == "D" diff --git a/src/tests/unit_tests/visualization/test_state_machine_viewer.py b/src/tests/unit_tests/visualization/test_state_machine_viewer.py new file mode 100644 index 000000000..5b25edf23 --- /dev/null +++ b/src/tests/unit_tests/visualization/test_state_machine_viewer.py @@ -0,0 +1,149 @@ +"""Tests for DecisionEvent recording and state machine visualization.""" + +import pytest + +from sherpa_ai.actions.empty import EmptyAction +from sherpa_ai.events import DecisionEvent, build_event +from sherpa_ai.memory.belief import Belief +from sherpa_ai.memory.state_machine import SherpaStateMachine +from sherpa_ai.visualization.state_machine_viewer import ( + DecisionRecord, + StateMachineViewer, +) + + +def test_decision_event_creation(): + event = DecisionEvent( + name="agent_1", chosen="A_to_B_1", alternatives=["A_to_B_2"], state="A", + ) + assert event.event_type == "decision" + assert event.chosen == "A_to_B_1" + assert event.alternatives == ["A_to_B_2"] + assert event.state == "A" + + +def test_decision_event_frozen_type(): + event = DecisionEvent(name="test", chosen="x", alternatives=[]) + assert event.event_type == "decision" + + +def test_build_event_decision(): + event = build_event( + "decision", "agent_1", + chosen="action_a", alternatives=["action_b", "action_c"], state="idle", + ) + assert isinstance(event, DecisionEvent) + assert event.chosen == "action_a" + assert event.alternatives == ["action_b", "action_c"] + assert event.state == "idle" + + +def test_belief_records_decision(): + belief = Belief() + belief.update_internal( + "decision", "agent_1", + chosen="search", alternatives=["deliberate", "summarize"], state="researching", + ) + assert len(belief.internal_events) == 1 + event = belief.internal_events[0] + assert isinstance(event, DecisionEvent) + assert event.chosen == "search" + + +def test_belief_get_by_type_decision(): + belief = Belief() + belief.update_internal("action_start", "agent_1", args={"query": "test"}) + belief.update_internal( + "decision", "agent_1", chosen="search", alternatives=["deliberate"], state="A", + ) + belief.update_internal("action_finish", "agent_1", outputs="result") + decisions = belief.get_by_type("decision") + assert len(decisions) == 1 + assert decisions[0].chosen == "search" + + +@pytest.fixture +def belief_with_sm(): + action_a = EmptyAction() + action_b = EmptyAction() + action_c = EmptyAction() + sm = SherpaStateMachine(states=["A", "B", "C"], initial="A") + sm.update_transition("A_to_B_1", "A", "B", action=action_a) + sm.update_transition("A_to_B_2", "A", "B", action=action_b) + sm.update_transition("B_to_C", "B", "C", action=action_c) + belief = Belief() + belief.state_machine = sm + belief.update_internal( + "decision", "test_agent", + chosen="A_to_B_1", alternatives=["A_to_B_2"], state="A", + ) + sm.A_to_B_1() + belief.update_internal( + "decision", "test_agent", chosen="B_to_C", alternatives=[], state="B", + ) + sm.B_to_C() + return belief + + +def test_extract_decisions(belief_with_sm): + viewer = StateMachineViewer(belief=belief_with_sm) + decisions = viewer.extract_decisions() + assert len(decisions) == 2 + assert decisions[0].step == 1 + assert decisions[0].state == "A" + assert decisions[0].chosen == "A_to_B_1" + assert decisions[0].alternatives == ["A_to_B_2"] + assert decisions[1].step == 2 + assert decisions[1].chosen == "B_to_C" + + +def test_extract_graph(belief_with_sm): + viewer = StateMachineViewer(belief=belief_with_sm) + graph = viewer.extract_graph() + assert set(graph["states"]) == {"A", "B", "C"} + triggers = {t["trigger"] for t in graph["transitions"]} + assert triggers == {"A_to_B_1", "A_to_B_2", "B_to_C"} + + +def test_extract_graph_no_state_machine(): + belief = Belief() + viewer = StateMachineViewer(belief=belief) + graph = viewer.extract_graph() + assert graph["states"] == [] + assert graph["transitions"] == [] + + +def test_render_html(belief_with_sm): + viewer = StateMachineViewer(belief=belief_with_sm) + html = viewer.render_html() + assert "" in html + assert "Sherpa State Machine Viewer" in html + assert ">A" in html + assert ">B" in html + assert ">C" in html + assert "A_to_B_1" in html + assert "B_to_C" in html + assert 'class="chosen"' in html + assert "A_to_B_2" in html + + +def test_render_to_file(belief_with_sm, tmp_path): + viewer = StateMachineViewer(belief=belief_with_sm) + output = tmp_path / "test_viz.html" + result_path = viewer.render(str(output)) + assert output.exists() + assert output.stat().st_size > 0 + content = output.read_text() + assert "Sherpa State Machine Viewer" in content + + +def test_render_empty_decisions(): + sm = SherpaStateMachine(states=["X", "Y"], initial="X") + sm.update_transition("go", "X", "Y", action=EmptyAction()) + belief = Belief() + belief.state_machine = sm + viewer = StateMachineViewer(belief=belief) + html = viewer.render_html() + assert "No decisions recorded yet." in html + assert ">X" in html + assert ">Y" in html \ No newline at end of file