Files
cleveragents-core/features/steps/conversation_state_management_steps.py
T
2026-06-18 09:49:41 -04:00

408 lines
15 KiB
Python

"""BDD step definitions for conversation state management feature."""
from __future__ import annotations
from behave import given, then, when
from behave.runner import Context
from unittest.mock import MagicMock
from cleveragents.domain.models.core.session import MessageRole, Session
from cleveragents.langgraph.bridge import RxPyLangGraphBridge
from cleveragents.langgraph.state import (
ConversationStateManager,
ExecutionState,
GraphState,
StateManager,
StateUpdateMode,
)
from cleveragents.reactive.stream_router import ReactiveStreamRouter
def _ensure(context: Context) -> None:
context.results = getattr(context, "results", {}) or {}
context.managers = getattr(context, "managers", {}) or {}
# ---------------------------------------------------------------------------
# ConversationStateManager scenarios
# ---------------------------------------------------------------------------
@given("a ConversationStateManager with no initial state")
def step_create_conversation_state_manager(context: Context) -> None:
_ensure(context)
context.managers["csm"] = ConversationStateManager()
@when('I append messages "hello" and "world" to the conversation history')
def step_append_hello_world(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.append_messages([{"role": "user", "content": "hello"}])
mgr.append_messages([{"role": "user", "content": "world"}])
@when("I reset the execution state")
def step_reset_execution_state(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.reset_execution_state()
@then("the conversation history should still contain 2 messages")
def step_assert_history_has_2_messages(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
assert len(mgr.get_full_history()) == 2, (
f"Expected 2 messages, got {len(mgr.get_full_history())}"
)
@then("the execution count should be 0 after reset")
def step_assert_execution_count_zero(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
state = mgr.get_state()
assert state.execution_count == 0, (
f"Expected execution_count=0, got {state.execution_count}"
)
@when('I update the state with current_node "agent_a" and execution_count 5')
def step_update_state_agent_a(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.update_state(
{"current_node": "agent_a"},
mode=StateUpdateMode.REPLACE,
)
# Perform 5 updates to increment execution_count
for _ in range(5):
mgr.update_state({}, mode=StateUpdateMode.REPLACE)
@when('I append a message "keep me" to the conversation history')
def step_append_keep_me(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.append_messages([{"role": "user", "content": "keep me"}])
@then('the conversation history should contain the message "keep me"')
def step_assert_history_contains_keep_me(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
history = mgr.get_full_history()
contents = [m.get("content") for m in history]
assert "keep me" in contents, f"'keep me' not found in history: {contents}"
@then("the current_node should be None after reset")
def step_assert_current_node_none(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
state = mgr.get_state()
assert state.current_node is None, (
f"Expected current_node=None, got {state.current_node!r}"
)
@when('I append messages "first" and "second" and "third" to the conversation history')
def step_append_three_messages(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
for content in ("first", "second", "third"):
mgr.append_messages([{"role": "user", "content": content}])
@then("get_full_history should return 3 messages")
def step_assert_full_history_3(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
history = mgr.get_full_history()
assert len(history) == 3, f"Expected 3 messages, got {len(history)}"
@when('I append a message "persistent" to the conversation history')
def step_append_persistent(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.append_messages([{"role": "user", "content": "persistent"}])
@then('the history property should contain the message "persistent"')
def step_assert_history_property_persistent(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
history = mgr.history
contents = [m.get("content") for m in history]
assert "persistent" in contents, (
f"'persistent' not found in history property: {contents}"
)
@when('I call update_state with messages containing "new message"')
def step_update_state_with_new_message(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.update_state(
{"messages": [{"role": "user", "content": "new message"}]},
mode=StateUpdateMode.APPEND,
)
@then('the conversation history should contain the message "new message"')
def step_assert_history_contains_new_message(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
history = mgr.get_full_history()
contents = [m.get("content") for m in history]
assert "new message" in contents, f"'new message' not found in history: {contents}"
# ---------------------------------------------------------------------------
# GraphState.to_graph_state scenario
# ---------------------------------------------------------------------------
@given("a GraphState with messages and metadata")
def step_create_graph_state_with_data(context: Context) -> None:
_ensure(context)
context.results["graph_state"] = GraphState(
messages=[{"role": "user", "content": "hi"}],
metadata={"key": "value"},
current_node="test_node",
execution_count=2,
)
@when("I call to_graph_state on the GraphState")
def step_call_to_graph_state(context: Context) -> None:
gs: GraphState = context.results["graph_state"]
context.results["graph_state_dict"] = gs.to_graph_state()
@then("the result should contain a conversation_history key")
def step_assert_conversation_history_key(context: Context) -> None:
result: dict = context.results["graph_state_dict"]
assert "conversation_history" in result, (
f"'conversation_history' key missing from: {list(result.keys())}"
)
assert isinstance(result["conversation_history"], list)
assert len(result["conversation_history"]) == 1
@then("the result should contain an execution_state key with current_node")
def step_assert_execution_state_key(context: Context) -> None:
result: dict = context.results["graph_state_dict"]
assert "execution_state" in result, (
f"'execution_state' key missing from: {list(result.keys())}"
)
exec_state = result["execution_state"]
assert exec_state["current_node"] == "test_node"
assert exec_state["execution_count"] == 2
# ---------------------------------------------------------------------------
# ExecutionState scenario
# ---------------------------------------------------------------------------
@given('an ExecutionState with current_node "node_x" and execution_count 3')
def step_create_execution_state(context: Context) -> None:
_ensure(context)
context.results["exec_state"] = ExecutionState(
current_node="node_x",
execution_count=3,
)
@then('the ExecutionState current_node should be "node_x"')
def step_assert_exec_state_current_node(context: Context) -> None:
es: ExecutionState = context.results["exec_state"]
assert es.current_node == "node_x", f"Expected 'node_x', got {es.current_node!r}"
@then("the ExecutionState execution_count should be 3")
def step_assert_exec_state_execution_count(context: Context) -> None:
es: ExecutionState = context.results["exec_state"]
assert es.execution_count == 3, f"Expected 3, got {es.execution_count}"
# ---------------------------------------------------------------------------
# StateManager backward-compat alias scenario
# ---------------------------------------------------------------------------
@given("a StateManager created via the backward-compat alias")
def step_create_state_manager_alias(context: Context) -> None:
_ensure(context)
context.managers["sm"] = StateManager()
@when('I append a message "compat" to the conversation history via StateManager')
def step_append_compat_via_state_manager(context: Context) -> None:
mgr: StateManager = context.managers["sm"]
mgr.append_messages([{"role": "user", "content": "compat"}])
@then('the StateManager history property should contain the message "compat"')
def step_assert_state_manager_history_compat(context: Context) -> None:
mgr: StateManager = context.managers["sm"]
history = mgr.history
contents = [m.get("content") for m in history]
assert "compat" in contents, (
f"'compat' not found in StateManager history: {contents}"
)
# ---------------------------------------------------------------------------
# MAX_HISTORY_SIZE eviction scenario
# ---------------------------------------------------------------------------
@when("I append 55 messages to the conversation history")
def step_append_55_messages(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
for i in range(55):
mgr.append_messages([{"role": "user", "content": f"msg_{i}"}])
@then("the conversation history should contain at most 50 messages")
def step_assert_history_at_most_50(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
history = mgr.get_full_history()
assert len(history) <= 50, f"Expected at most 50 messages, got {len(history)}"
@then("the most recent messages should be retained")
def step_assert_most_recent_retained(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
history = mgr.get_full_history()
# The last message appended was msg_54
assert history[-1]["content"] == "msg_54", (
f"Expected last message to be 'msg_54', got {history[-1]['content']!r}"
)
# ---------------------------------------------------------------------------
# Full reset scenario
# ---------------------------------------------------------------------------
@when('I append a message "before reset" to the conversation history')
def step_append_before_reset(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.append_messages([{"role": "user", "content": "before reset"}])
@when('I update the state with current_node "some_node" and execution_count 5')
def step_update_state_some_node(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.update_state({"current_node": "some_node"}, mode=StateUpdateMode.REPLACE)
for _ in range(5):
mgr.update_state({}, mode=StateUpdateMode.REPLACE)
@when("I perform a full reset")
def step_perform_full_reset(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
mgr.reset()
@then("the conversation history should be empty after full reset")
def step_assert_history_empty_after_full_reset(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
history = mgr.get_full_history()
assert len(history) == 0, f"Expected empty history, got {len(history)} messages"
@then("the execution count should be 0 after full reset")
def step_assert_execution_count_zero_after_full_reset(context: Context) -> None:
mgr: ConversationStateManager = context.managers["csm"]
state = mgr.get_state()
assert state.execution_count == 0, (
f"Expected execution_count=0, got {state.execution_count}"
)
# ---------------------------------------------------------------------------
# Bridge default update mode scenario
# ---------------------------------------------------------------------------
@given("a RxPyLangGraphBridge instance")
def step_create_bridge(context: Context) -> None:
_ensure(context)
scheduler = MagicMock()
router = ReactiveStreamRouter(scheduler=scheduler)
context.results["bridge"] = RxPyLangGraphBridge(stream_router=router)
@then("the default update mode should be APPEND")
def step_assert_bridge_default_mode_append(context: Context) -> None:
bridge: RxPyLangGraphBridge = context.results["bridge"]
assert bridge._DEFAULT_UPDATE_MODE == StateUpdateMode.APPEND, (
f"Expected APPEND, got {bridge._DEFAULT_UPDATE_MODE}"
)
# ---------------------------------------------------------------------------
# Session scenarios
# ---------------------------------------------------------------------------
def _make_session() -> Session:
from ulid import ULID
return Session(session_id=str(ULID()))
@given("a Session with two messages appended")
def step_session_with_two_messages(context: Context) -> None:
_ensure(context)
session = _make_session()
session.append_message(MessageRole.USER, "first message")
session.append_message(MessageRole.ASSISTANT, "first response")
context.results["session"] = session
@then("the full_history property should return 2 messages")
def step_assert_full_history_2(context: Context) -> None:
session: Session = context.results["session"]
full_history = session.full_history
assert len(full_history) == 2, (
f"Expected 2 messages in full_history, got {len(full_history)}"
)
@given("a Session with three messages appended")
def step_session_with_three_messages(context: Context) -> None:
_ensure(context)
session = _make_session()
session.append_message(MessageRole.USER, "msg one")
session.append_message(MessageRole.ASSISTANT, "msg two")
session.append_message(MessageRole.USER, "msg three")
context.results["session"] = session
@when("I call get_messages with no limit")
def step_call_get_messages_no_limit(context: Context) -> None:
session: Session = context.results["session"]
context.results["messages"] = session.get_messages(limit=None)
@then("all 3 messages should be returned")
def step_assert_all_3_messages(context: Context) -> None:
messages = context.results["messages"]
assert len(messages) == 3, f"Expected 3 messages, got {len(messages)}"
@given("a Session with one existing message")
def step_session_with_one_message(context: Context) -> None:
_ensure(context)
session = _make_session()
session.append_message(MessageRole.USER, "first")
context.results["session"] = session
@when("I append a second message to the session")
def step_append_second_message(context: Context) -> None:
session: Session = context.results["session"]
session.append_message(MessageRole.ASSISTANT, "second")
@then("the session should have 2 messages in total")
def step_assert_session_has_2_messages(context: Context) -> None:
session: Session = context.results["session"]
assert session.message_count == 2, (
f"Expected 2 messages, got {session.message_count}"
)