408 lines
15 KiB
Python
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}"
|
|
)
|