474 lines
16 KiB
Python
474 lines
16 KiB
Python
"""Step definitions for LangGraph integration tests."""
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch
|
|
|
|
import yaml
|
|
from behave import given
|
|
from behave import then
|
|
from behave import when
|
|
from rx.subject import Subject
|
|
|
|
from cleveragents.core.application import ReactiveCleverAgentsApp
|
|
from cleveragents.langgraph.graph import GraphConfig
|
|
from cleveragents.langgraph.graph import LangGraph
|
|
from cleveragents.langgraph.nodes import Edge
|
|
from cleveragents.langgraph.nodes import NodeConfig
|
|
from cleveragents.langgraph.nodes import NodeType
|
|
from cleveragents.langgraph.state import GraphState
|
|
|
|
|
|
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("cleveragents.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 = loop.run_until_complete(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 = loop.run_until_complete(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("cleveragents.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 cleveragents.reactive.stream_router import StreamConfig
|
|
from cleveragents.reactive.stream_router import 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."""
|
|
# Placeholder for visualization test
|
|
pass
|
|
|
|
|
|
@then("the diagram should show all edges with conditions if present")
|
|
def step_check_mermaid_edges(context):
|
|
"""Check mermaid diagram has edges."""
|
|
# Placeholder for visualization test
|
|
pass
|
|
|
|
|
|
@then("the diagram should show subgraph boundaries if applicable")
|
|
def step_check_mermaid_subgraphs(context):
|
|
"""Check mermaid diagram has subgraphs."""
|
|
# Placeholder for visualization test
|
|
pass
|