cleveractors-core/features/steps/langgraph_steps.py

"""Step definitions for LangGraph integration tests."""

import asyncio
import json
from unittest.mock import MagicMock, patch

import yaml
from behave import given, then, when
from rx.subject import Subject

from cleveractors.core.application import ReactiveCleverAgentsApp
from cleveractors.langgraph.graph import GraphConfig, LangGraph
from cleveractors.langgraph.nodes import Edge, NodeConfig, NodeType


def create_mock_langgraph(name="test_graph", nodes=None, edges=None, config_dict=None):
    """Create a mock LangGraph for testing."""
    mock_graph = MagicMock()
    mock_graph.name = name
    mock_graph.config = MagicMock()
    mock_graph.config.entry_point = "start"

    # Set defaults
    mock_graph.config.checkpointing = False
    mock_graph.config.enable_time_travel = False
    mock_graph.config.parallel_execution = False

    # Override with configuration if provided
    if config_dict:
        mock_graph.config.checkpointing = config_dict.get("checkpointing", False)
        mock_graph.config.enable_time_travel = config_dict.get(
            "enable_time_travel", False
        )
        mock_graph.config.parallel_execution = config_dict.get(
            "parallel_execution", False
        )
    mock_graph.config.nodes = nodes or {}
    # Create proper edge objects with conditions
    mock_edges = []
    if edges:
        for edge in edges:
            mock_edge = MagicMock()
            mock_edge.source = edge.get("source")
            mock_edge.target = edge.get("target")
            mock_edge.condition = edge.get("condition")
            mock_edges.append(mock_edge)
    mock_graph.config.edges = mock_edges
    # Ensure nodes include start, end, and any custom nodes
    graph_nodes = {}
    if nodes:
        for node_name, node_config in nodes.items():
            mock_node = MagicMock()
            mock_node.config = MagicMock()
            mock_node.config.type = NodeType.FUNCTION  # default
            if node_config.get("type") == "agent":
                mock_node.config.type = NodeType.AGENT
                mock_node.config.agent = node_config.get("agent")
            elif node_config.get("type") == "conditional":
                mock_node.config.type = NodeType.CONDITIONAL
            graph_nodes[node_name] = mock_node
    # Always add start and end nodes
    graph_nodes.update({"start": MagicMock(), "end": MagicMock()})
    mock_graph.nodes = graph_nodes

    # Mock state manager
    mock_graph.state_manager = MagicMock()
    mock_graph.state_manager.get_state = MagicMock()
    mock_graph.state_manager.update_state = MagicMock()
    mock_graph.state_manager.get_state_observable = MagicMock(return_value=Subject())

    # Mock execute method
    async def mock_execute(input_data):
        state = MagicMock()
        state.messages = [{"role": "assistant", "content": "Mock response"}]
        state.to_dict = MagicMock(return_value={"messages": state.messages})
        return state

    mock_graph.execute = mock_execute
    mock_graph.get_execution_history = MagicMock(return_value=[])
    mock_graph.visualize = MagicMock(return_value="Mock visualization")

    return mock_graph


@given("the CleverAgents reactive system is available with LangGraph support")
def step_impl(context):
    """Initialize the reactive system with LangGraph support."""
    # Clear any existing state to prevent pollution from other tests
    if hasattr(context, "app"):
        context.app = None
    if hasattr(context, "graphs"):
        context.graphs.clear()
    if hasattr(context, "config_dict"):
        context.config_dict.clear()

    # Create fresh instances
    context.app = ReactiveCleverAgentsApp(verbose=True)
    context.graphs = {}
    context.config_dict = {}


@given("I have a LangGraph configuration")
def step_impl(context):
    """Parse LangGraph configuration from text."""
    context.graph_config = json.loads(context.text)


@given("I have agents configured")
def step_impl(context):
    """Parse agent configuration."""
    context.config_dict.update(yaml.safe_load(context.text))


@given("I have a hybrid pipeline configuration")
def step_impl(context):
    """Parse hybrid pipeline configuration."""
    context.config_dict.update(yaml.safe_load(context.text))


@when("I create the graph")
def step_impl(context):
    """Create a LangGraph from configuration."""
    try:
        # Ensure we have valid graph config
        if not hasattr(context, "graph_config") or not context.graph_config:
            raise ValueError("No graph configuration provided")

        # Ensure we have a valid app instance
        if not hasattr(context, "app") or not context.app:
            raise ValueError("No ReactiveCleverAgentsApp instance available")

        graph_name = context.graph_config.get("name", "test_graph")

        with patch("cleveractors.langgraph.bridge.LangGraph") as mock_langgraph_class:
            # Create mock with configuration details
            mock_graph = create_mock_langgraph(
                name=graph_name,
                nodes=context.graph_config.get("nodes", {}),
                edges=context.graph_config.get("edges", []),
                config_dict=context.graph_config,
            )
            mock_langgraph_class.return_value = mock_graph

            # Create the graph through the bridge
            graph = context.app.langgraph_bridge.create_graph_from_config(
                context.graph_config
            )
            context.graphs[graph.name] = graph
            context.current_graph = graph
            context.graph_created = True
    except Exception as e:
        context.graph_created = False
        context.error = str(e)


@when("I create the pipeline")
def step_impl(context):
    """Create a hybrid pipeline."""
    try:
        # Load configuration into app
        context.app.config = context.app.config_parser._build_reactive_config(
            context.config_dict
        )

        # Set up pipelines
        context.app._setup_pipelines()
        context.pipeline_created = True
    except Exception as e:
        context.pipeline_created = False
        context.error = str(e)


@when("I execute the graph with messages")
def step_impl(context):
    """Execute a graph with test messages."""
    graph = context.current_graph
    messages = [
        {"role": "user", "content": "Test message 1"},
        {"role": "user", "content": "Test message 2"},
    ]

    # Execute graph
    loop = asyncio.get_event_loop()
    context.final_state = asyncio.run(graph.execute({"messages": messages}))


@when("I execute the graph")
def step_impl(context):
    """Execute a graph without specific input."""
    graph = context.current_graph

    # Execute graph
    loop = asyncio.get_event_loop()
    context.final_state = asyncio.run(graph.execute())
    context.execution_history = graph.get_execution_history()


@when("I send a message to the stream")
def step_impl(context):
    """Send a test message to a stream."""
    context.app.stream_router.send_message(
        context.stream_config["name"], "Test message for graph processing"
    )


@when("I send state updates through the stream")
def step_impl(context):
    """Send state updates through a stream."""
    updates = {
        "metadata": {"updated": True, "counter": 1},
        "messages": [{"role": "system", "content": "State updated"}],
    }
    context.app.stream_router.send_message(context.stream_config["name"], updates)


@when("I request visualization in mermaid format")
def step_impl(context):
    """Request graph visualization."""
    if hasattr(context, "current_graph"):
        context.visualization = context.current_graph.visualize("mermaid")
    else:
        # Create a complex graph for visualization
        config = GraphConfig(
            name="complex_graph",
            nodes={
                "analyze": NodeConfig(name="analyze", type=NodeType.AGENT),
                "validate": NodeConfig(name="validate", type=NodeType.FUNCTION),
                "route": NodeConfig(name="route", type=NodeType.CONDITIONAL),
            },
            edges=[
                Edge("start", "analyze"),
                Edge("analyze", "validate"),
                Edge("validate", "route"),
                Edge("route", "end", condition={"type": "always"}),
            ],
        )
        graph = LangGraph(config)
        context.visualization = graph.visualize("mermaid")


@then("the graph should be created successfully")
def step_impl(context):
    """Verify graph creation."""
    assert context.graph_created, (
        f"Graph creation failed: {getattr(context, 'error', 'Unknown error')}"
    )
    assert context.current_graph is not None


@then("the graph should have {count:d} custom node plus start and end nodes")
def step_impl(context, count):
    """Verify node count."""
    graph = context.current_graph
    # Total nodes = custom nodes + start + end
    assert len(graph.nodes) == count + 2


@then("the graph should connect to the analyzer agent")
def step_impl(context):
    """Verify agent connection."""
    graph = context.current_graph
    analyze_node = graph.nodes.get("analyze")
    assert analyze_node is not None
    assert analyze_node.config.type == NodeType.AGENT
    assert analyze_node.config.agent == "analyzer"


@then("the graph should support conditional routing")
def step_impl(context):
    """Verify conditional routing support."""
    graph = context.current_graph

    # Check for conditional node
    check_node = graph.nodes.get("check")
    assert check_node is not None
    assert check_node.config.type == NodeType.CONDITIONAL

    # Check for conditional edges
    conditional_edges = [e for e in graph.config.edges if e.condition is not None]
    assert len(conditional_edges) > 0


@then("the graph state should be persisted")
def step_impl(context):
    """Verify state persistence."""
    graph = context.current_graph
    assert graph.config.checkpointing is True

    # Check if checkpoint exists
    latest_checkpoint = graph.state_manager.get_latest_checkpoint()
    # In test environment, checkpoint_dir might not be set
    # so we just verify the checkpointing flag


@then("I should be able to time travel to previous states")
def step_impl(context):
    """Verify time travel capability."""
    graph = context.current_graph
    assert graph.config.enable_time_travel is True

    # Try time travel
    previous_state = graph.state_manager.time_travel(1)
    # In test, we might not have history, so just verify the feature is enabled


@then("the pipeline should combine RxPy streams and LangGraph execution")
def step_impl(context):
    """Verify hybrid pipeline creation."""
    assert context.pipeline_created, (
        f"Pipeline creation failed: {getattr(context, 'error', 'Unknown error')}"
    )

    # Check that both streams and graphs exist
    assert "input_processor" in context.app.stream_router.streams
    assert "processing_graph" in context.app.langgraph_bridge.graphs


@then("task1 and task2 should execute in parallel")
def step_impl(context):
    """Verify parallel execution."""
    # Check execution history to verify parallel execution
    history = context.execution_history

    # Find indices of task1 and task2
    task1_idx = history.index("task1") if "task1" in history else -1
    task2_idx = history.index("task2") if "task2" in history else -1
    combine_idx = history.index("combine") if "combine" in history else -1

    # Both tasks should be executed before combine
    assert task1_idx < combine_idx
    assert task2_idx < combine_idx

    # Verify parallel execution flag
    assert context.current_graph.config.parallel_execution is True


@then("the message should be processed by the LangGraph")
def step_impl(context):
    """Verify graph processing through stream."""
    # In a real test, we would check the output
    # For now, verify the stream was created with graph operator
    assert "graph_execute" in str(context.stream_config["operators"])


@then("the graph state should be updated accordingly")
def step_impl(context):
    """Verify state updates."""
    # In a real test, we would check the actual state
    # For now, verify the stream was created with state update operator
    assert "state_update" in str(context.stream_config["operators"])


@then("I should get a mermaid diagram showing")
def step_impl(context):
    """Verify visualization output."""
    viz = context.visualization
    assert viz.startswith("graph TD")

    # Check for expected elements in the visualization
    for row in context.table:
        element = row["element"]
        if element == "All nodes with appropriate shapes":
            assert "[" in viz  # Check for node shapes
        elif element == "All edges with conditions if present":
            assert "-->" in viz  # Check for edges
        elif element == "Subgraph boundaries if applicable":
            # Subgraphs would contain "subgraph" keyword
            pass  # Optional in simple graphs


@given("I have created the {graph_name}")
def step_impl(context, graph_name):
    """Create a named graph for testing."""
    config = {
        "name": graph_name,
        "nodes": {"process": {"type": "function", "function": "summarize"}},
        "edges": [
            {"source": "start", "target": "process"},
            {"source": "process", "target": "end"},
        ],
    }

    if graph_name == "stateful_graph":
        config["checkpointing"] = True
        config["enable_time_travel"] = True

    with patch("cleveractors.langgraph.bridge.LangGraph") as mock_langgraph_class:
        mock_graph = create_mock_langgraph(
            name=graph_name,
            nodes=config.get("nodes", {}),
            edges=config.get("edges", []),
            config_dict=config,
        )
        mock_langgraph_class.return_value = mock_graph

        graph = context.app.langgraph_bridge.create_graph_from_config(config)
        context.graphs[graph_name] = graph


@given("I have a stream configuration with LangGraph operator")
def step_impl(context):
    """Parse stream configuration with LangGraph operator."""
    context.stream_config = json.loads(context.text)


@given("I have a stream configuration with state update operator")
def step_impl(context):
    """Parse stream configuration with state update operator."""
    context.stream_config = json.loads(context.text)


@when("I create the stream")
def step_impl(context):
    """Create a stream from configuration."""
    from cleveractors.reactive.stream_router import StreamConfig, StreamType

    # Handle both dict and StreamConfig object
    if isinstance(context.stream_config, dict):
        config = StreamConfig(
            name=context.stream_config["name"],
            type=StreamType(context.stream_config.get("type", "cold")),
            operators=context.stream_config.get("operators", []),
            publications=context.stream_config.get("publications", []),
        )
    else:
        config = context.stream_config

    # Handle different contexts - reactive_steps vs langgraph_steps
    if hasattr(context, "app") and context.app:
        # LangGraph context
        context.app.stream_router._langgraph_bridge = context.app.langgraph_bridge
        context.app.stream_router.create_stream(config)
    elif hasattr(context, "stream_router"):
        # Reactive context
        context.stream_router.create_stream(config)
    else:
        raise ValueError("No stream router available in context")


@given("I have a complex LangGraph configured")
def step_complex_langgraph(context):
    """Create a complex LangGraph for testing."""
    context.graph_created = True


@then("I should get a mermaid diagram showing all nodes with appropriate shapes")
def step_check_mermaid_nodes(context):
    """Check mermaid diagram has nodes."""
    assert context.visualization is not None, "Visualization should be produced"
    assert isinstance(context.visualization, str), (
        f"Visualization should be a string, got {type(context.visualization)}"
    )
    assert "graph TD" in context.visualization or "graph LR" in context.visualization, (
        "Mermaid diagram should start with graph declaration"
    )
    assert "analyze" in context.visualization, "Should contain analyze node"
    assert "validate" in context.visualization, "Should contain validate node"
    assert "route" in context.visualization, "Should contain route node"


@then("the diagram should show all edges with conditions if present")
def step_check_mermaid_edges(context):
    """Check mermaid diagram has edges."""
    assert context.visualization is not None, "Visualization should be produced"
    viz = context.visualization
    assert "-->" in viz or "->" in viz, (
        f"Mermaid diagram should contain edge arrows, got: {viz[:100]}..."
    )


@then("the diagram should show subgraph boundaries if applicable")
def step_check_mermaid_subgraphs(context):
    """Check mermaid diagram has subgraphs or valid structure."""
    assert context.visualization is not None, "Visualization should be produced"
    assert len(context.visualization) > 20, (
        f"Visualization should contain meaningful content, got len {len(context.visualization)}"
    )