forked from HAL9000/cleveragents-core
844 lines
30 KiB
Python
844 lines
30 KiB
Python
"""Behave steps covering memory_service module."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import shutil
|
|
import tempfile
|
|
import time
|
|
from datetime import UTC
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
from behave import given, then, when
|
|
from langchain_core.chat_history import InMemoryChatMessageHistory
|
|
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
|
|
|
|
from cleveragents.application.services.memory_service import (
|
|
ConversationBufferMemoryAdapter,
|
|
Entity,
|
|
EntityStore,
|
|
EntityType,
|
|
MemoryService,
|
|
)
|
|
|
|
|
|
@given("a conversation adapter configured to return text values")
|
|
def step_adapter_returning_text(context: Any) -> None:
|
|
"""Create an adapter that returns buffer strings."""
|
|
|
|
history = InMemoryChatMessageHistory()
|
|
history.add_user_message("Hello from user")
|
|
history.add_ai_message("Greetings from AI")
|
|
context.adapter = ConversationBufferMemoryAdapter(
|
|
history,
|
|
return_messages=False,
|
|
)
|
|
context.adapter_keys = context.adapter.memory_variables
|
|
|
|
|
|
@when("I load memory variables from the adapter")
|
|
def step_load_memory_variables(context: Any) -> None:
|
|
"""Load memory payload from the adapter."""
|
|
|
|
context.loaded_memory = context.adapter.load_memory_variables()
|
|
|
|
|
|
@then("the adapter provides a string buffer representation")
|
|
def step_validate_string_payload(context: Any) -> None:
|
|
"""Ensure adapter returns buffer string when configured."""
|
|
|
|
assert context.adapter.memory_key in context.loaded_memory
|
|
payload = context.loaded_memory[context.adapter.memory_key]
|
|
assert isinstance(payload, str)
|
|
assert payload
|
|
assert "Hello from user" in payload or "Greetings from AI" in payload
|
|
assert context.adapter_keys == [context.adapter.memory_key]
|
|
|
|
|
|
@given("a conversation adapter with prune tracking")
|
|
def step_adapter_with_prune(context: Any) -> None:
|
|
"""Create adapter that records prune callback invocations."""
|
|
|
|
history = InMemoryChatMessageHistory()
|
|
context.prune_calls = 0
|
|
|
|
def record_prune() -> None:
|
|
context.prune_calls += 1
|
|
|
|
context.adapter = ConversationBufferMemoryAdapter(
|
|
history, prune_callback=record_prune
|
|
)
|
|
context.adapter_keys = context.adapter.memory_variables
|
|
|
|
|
|
@when("I save nested user and ai context through the adapter")
|
|
def step_save_nested_context(context: Any) -> None:
|
|
"""Store nested values to exercise adapter recursion."""
|
|
|
|
inputs = {
|
|
"input": [
|
|
HumanMessage(content="first human"),
|
|
"plain user",
|
|
["nested user", HumanMessage(content="second human")],
|
|
None,
|
|
]
|
|
}
|
|
outputs = {
|
|
"output": [
|
|
AIMessage(content="first ai"),
|
|
"plain ai",
|
|
[AIMessage(content="second ai"), "another ai"],
|
|
None,
|
|
]
|
|
}
|
|
context.adapter.save_context(inputs, outputs)
|
|
context.saved_snapshot = [
|
|
msg.content for msg in context.adapter._message_history.messages
|
|
]
|
|
|
|
|
|
@when("I invoke the adapter async operations")
|
|
def step_async_adapter_ops(context: Any) -> None:
|
|
"""Exercise async adapter APIs."""
|
|
|
|
async def run_ops() -> None:
|
|
context.async_loaded = await context.adapter.aload_memory_variables({})
|
|
await context.adapter.asave_context(
|
|
{"input": [HumanMessage(content="async human"), ["async text"]]},
|
|
{"output": [AIMessage(content="async ai"), ["async reply"]]},
|
|
)
|
|
context.post_async_snapshot = [
|
|
msg.content for msg in context.adapter._message_history.messages
|
|
]
|
|
await context.adapter.aclear()
|
|
|
|
asyncio.run(run_ops())
|
|
|
|
|
|
@then("the adapter normalizes all values and prunes history three times")
|
|
def step_verify_prune_and_storage(context: Any) -> None:
|
|
"""Validate adapter processed values and triggered pruning."""
|
|
|
|
assert context.prune_calls == 3
|
|
assert len(context.saved_snapshot) >= 4
|
|
assert any("plain user" in content for content in context.saved_snapshot)
|
|
async_payload = context.async_loaded[context.adapter.memory_key]
|
|
assert isinstance(async_payload, list)
|
|
assert any(isinstance(msg, BaseMessage) for msg in async_payload)
|
|
assert len(context.post_async_snapshot) >= 2
|
|
assert not context.adapter._message_history.messages
|
|
|
|
|
|
@given("a memory service limited to {count:d} messages")
|
|
def step_memory_service_with_limit(context: Any, count: int) -> None:
|
|
"""Create memory service with max message window."""
|
|
|
|
context.memory_service = MemoryService(
|
|
session_id="limited-session", max_messages=count
|
|
)
|
|
|
|
|
|
@when("I add five interactions via save context")
|
|
def step_add_multiple_interactions(context: Any) -> None:
|
|
"""Add several interactions to trigger pruning."""
|
|
|
|
for idx in range(5):
|
|
context.memory_service.save_context(
|
|
{"input": f"user {idx}"},
|
|
{"output": f"ai {idx}"},
|
|
)
|
|
context.final_messages = context.memory_service.get_messages()
|
|
|
|
|
|
@when("I ask for the two most recent messages")
|
|
def step_recent_two(context: Any) -> None:
|
|
"""Fetch limited recent messages."""
|
|
|
|
context.recent_two = context.memory_service.get_recent_messages(2)
|
|
|
|
|
|
@when("I ask for more messages than available")
|
|
def step_recent_all(context: Any) -> None:
|
|
"""Fetch more messages than stored to exercise branch."""
|
|
|
|
context.recent_all = context.memory_service.get_recent_messages(10)
|
|
|
|
|
|
@then("the service keeps only the three latest messages")
|
|
def step_verify_pruned_messages(context: Any) -> None:
|
|
"""Ensure pruning retained the configured window."""
|
|
|
|
assert len(context.final_messages) == 3
|
|
contents = [msg.content for msg in context.final_messages]
|
|
assert contents[-1] == "ai 4"
|
|
assert "user 4" in contents[-2]
|
|
|
|
|
|
@then("the recent message requests respect their sizes")
|
|
def step_verify_recent_requests(context: Any) -> None:
|
|
"""Validate recent message helpers."""
|
|
|
|
assert len(context.recent_two) == 2
|
|
assert len(context.recent_all) == len(context.final_messages)
|
|
assert [msg.content for msg in context.recent_two] == ["user 4", "ai 4"]
|
|
|
|
|
|
@then("the service summary honors a 20 character limit")
|
|
def step_verify_summary_limit(context: Any) -> None:
|
|
"""Confirm summary respects the provided character cap."""
|
|
|
|
summary = context.memory_service.get_summary(max_chars=20)
|
|
context.summary = summary
|
|
lines = summary.splitlines()
|
|
assert 1 <= len(lines) <= 2
|
|
assert lines[-1].endswith("ai 4")
|
|
|
|
|
|
@then("memory variables expose history and counts")
|
|
def step_verify_memory_variables(context: Any) -> None:
|
|
"""Check memory variables include history, summary, and count."""
|
|
|
|
variables = context.memory_service.get_memory_variables()
|
|
assert variables["message_count"] == len(context.final_messages)
|
|
assert variables["chat_history"] == context.final_messages
|
|
assert isinstance(variables["chat_summary"], str)
|
|
assert context.final_messages[-1].content in variables["chat_summary"]
|
|
if context.summary:
|
|
assert (
|
|
context.summary.splitlines()[-1].split(":", 1)[-1].strip()
|
|
in variables["chat_summary"]
|
|
)
|
|
|
|
|
|
@given("a memory service with default settings")
|
|
def step_memory_service_default(context: Any) -> None:
|
|
"""Create default memory service without max limit."""
|
|
|
|
context.memory_service = MemoryService(session_id="default-session")
|
|
|
|
|
|
@when("I add direct message objects to the service")
|
|
def step_add_direct_messages(context: Any) -> None:
|
|
"""Add BaseMessage instances directly and capture summary/token count."""
|
|
|
|
context.memory_service.add_message(
|
|
HumanMessage(content="Alpha conversation segment")
|
|
)
|
|
context.memory_service.add_message(AIMessage(content="Beta response text"))
|
|
context.summary_before_clear = context.memory_service.get_summary()
|
|
context.token_count_before_clear = context.memory_service.get_token_count()
|
|
context.messages_before_clear = context.memory_service.get_messages().copy()
|
|
|
|
|
|
@when("I inspect the default conversation adapter")
|
|
def step_inspect_default_adapter(context: Any) -> None:
|
|
"""Grab default adapter payload."""
|
|
|
|
context.default_adapter = context.memory_service.conversation_memory
|
|
context.default_payload = context.default_adapter.load_memory_variables()
|
|
|
|
|
|
@when("I reset the max message limit to 1")
|
|
def step_reset_max_messages(context: Any) -> None:
|
|
"""Reduce max message count and trigger immediate pruning."""
|
|
|
|
context.memory_service.set_max_messages(1)
|
|
context.post_set_messages = list(context.memory_service.get_messages())
|
|
|
|
|
|
@when("I save a new interaction to enforce pruning")
|
|
def step_save_interaction_after_limit(context: Any) -> None:
|
|
"""Save interaction that should prune to configured limit."""
|
|
|
|
context.memory_service.save_context(
|
|
{"input": "latest user"},
|
|
{"output": "latest ai"},
|
|
)
|
|
context.pruned_messages = context.memory_service.get_messages()
|
|
|
|
|
|
@when("I clear the service memory")
|
|
def step_clear_service(context: Any) -> None:
|
|
"""Clear the memory service history."""
|
|
|
|
context.memory_service.clear()
|
|
|
|
|
|
@then("the service computed token counts and summaries before clearing")
|
|
def step_validate_pre_clear_metrics(context: Any) -> None:
|
|
"""Ensure summary and token counts were recorded before clearing."""
|
|
|
|
assert context.summary_before_clear
|
|
assert context.token_count_before_clear > 0
|
|
payload = context.default_payload[context.default_adapter.memory_key]
|
|
assert isinstance(payload, list)
|
|
assert len(context.post_set_messages) == 1
|
|
|
|
|
|
@then("the service history is empty after clearing")
|
|
def step_validate_cleared_history(context: Any) -> None:
|
|
"""History should be empty after clear."""
|
|
|
|
assert not context.memory_service.get_messages()
|
|
|
|
|
|
@then("a custom conversation adapter returns text values")
|
|
def step_custom_adapter_returns_text(context: Any) -> None:
|
|
"""Create custom adapter and ensure it returns string payloads."""
|
|
|
|
custom_adapter = context.memory_service.create_conversation_memory(
|
|
memory_key="custom_history",
|
|
input_key="prompt",
|
|
output_key="reply",
|
|
return_messages=False,
|
|
)
|
|
payload = custom_adapter.load_memory_variables()
|
|
assert payload["custom_history"] == ""
|
|
|
|
|
|
@given("a patched SQL chat history for memory service")
|
|
def step_patch_sql_history(context: Any) -> None:
|
|
"""Patch SQLChatMessageHistory to observe construction."""
|
|
|
|
context.sql_history_calls = []
|
|
context.sql_history_instance = InMemoryChatMessageHistory()
|
|
context.sql_temp_dir = tempfile.mkdtemp()
|
|
|
|
def factory(*args: Any, **kwargs: Any) -> InMemoryChatMessageHistory:
|
|
context.sql_history_calls.append((args, kwargs))
|
|
return context.sql_history_instance
|
|
|
|
sql_patch = patch(
|
|
"cleveragents.application.services.memory_service.SQLChatMessageHistory",
|
|
side_effect=factory,
|
|
)
|
|
context.sql_patch = sql_patch
|
|
sql_patch.start()
|
|
if hasattr(context, "_cleanup_handlers"):
|
|
context._cleanup_handlers.append(sql_patch.stop)
|
|
context._cleanup_handlers.append(
|
|
lambda: shutil.rmtree(context.sql_temp_dir, ignore_errors=True)
|
|
)
|
|
|
|
|
|
@when("I create a memory service with a connection string")
|
|
def step_create_sql_memory_service(context: Any) -> None:
|
|
"""Instantiate memory service pointing at SQL backend."""
|
|
|
|
db_path = Path(context.sql_temp_dir) / "memory.db"
|
|
context.connection_string = f"sqlite:///{db_path}"
|
|
context.sql_service = MemoryService(
|
|
session_id="sql-session",
|
|
connection_string=context.connection_string,
|
|
)
|
|
|
|
|
|
@then("the SQL history constructor receives the session and connection")
|
|
def step_validate_sql_constructor(context: Any) -> None:
|
|
"""Verify the patched SQL history constructor was invoked correctly."""
|
|
|
|
assert context.sql_history_calls
|
|
_, kwargs = context.sql_history_calls[0]
|
|
assert kwargs["session_id"] == "sql-session"
|
|
assert kwargs["connection_string"] == context.connection_string
|
|
assert kwargs.get("table_name") == "message_history"
|
|
|
|
|
|
@then("stored messages flow through the patched SQL history")
|
|
def step_validate_sql_storage(context: Any) -> None:
|
|
"""Ensure messages saved through the service reach the patched history."""
|
|
|
|
context.sql_service.save_context(
|
|
{"input": "persisted human"},
|
|
{"output": "persisted ai"},
|
|
)
|
|
messages = context.sql_history_instance.messages
|
|
assert len(messages) == 2
|
|
assert isinstance(messages[0], HumanMessage)
|
|
assert isinstance(messages[1], AIMessage)
|
|
assert messages[0].content == "persisted human"
|
|
assert messages[1].content == "persisted ai"
|
|
|
|
|
|
# ==================== EntityMemory Steps ====================
|
|
|
|
|
|
@given("a memory service with entity tracking enabled")
|
|
def step_memory_service_with_entities(context: Any) -> None:
|
|
"""Create a memory service for entity tracking tests."""
|
|
context.memory_service = MemoryService(session_id="entity-test-session")
|
|
|
|
|
|
@when('I track a project entity named "{name}"')
|
|
def step_track_project_entity(context: Any, name: str) -> None:
|
|
"""Track a project entity."""
|
|
context.tracked_entity = context.memory_service.track_entity(
|
|
name, EntityType.PROJECT
|
|
)
|
|
|
|
|
|
@then("the entity should be stored successfully")
|
|
def step_entity_stored(context: Any) -> None:
|
|
"""Verify entity was stored."""
|
|
assert context.tracked_entity is not None
|
|
assert isinstance(context.tracked_entity, Entity)
|
|
|
|
|
|
@then('the entity should have type "{entity_type}"')
|
|
def step_entity_has_type(context: Any, entity_type: str) -> None:
|
|
"""Verify entity type."""
|
|
assert context.tracked_entity.entity_type.value == entity_type
|
|
|
|
|
|
@then("the entity mention count should be {count:d}")
|
|
def step_entity_mention_count(context: Any, count: int) -> None:
|
|
"""Verify entity mention count."""
|
|
assert context.tracked_entity.mention_count == count
|
|
|
|
|
|
@when('I track a plan entity named "{name}" with metadata')
|
|
def step_track_plan_with_metadata(context: Any, name: str) -> None:
|
|
"""Track a plan entity with metadata."""
|
|
context.entity_metadata = {"status": "active", "priority": "high"}
|
|
context.tracked_entity = context.memory_service.track_entity(
|
|
name, EntityType.PLAN, context.entity_metadata
|
|
)
|
|
|
|
|
|
@then("the entity metadata should include the provided values")
|
|
def step_entity_metadata_includes_values(context: Any) -> None:
|
|
"""Verify metadata was stored."""
|
|
assert context.tracked_entity.metadata.get("status") == "active"
|
|
assert context.tracked_entity.metadata.get("priority") == "high"
|
|
|
|
|
|
@then("the entity should be retrievable by name and type")
|
|
def step_entity_retrievable(context: Any) -> None:
|
|
"""Verify entity can be retrieved."""
|
|
retrieved = context.memory_service.get_entity(
|
|
context.tracked_entity.name, context.tracked_entity.entity_type
|
|
)
|
|
assert retrieved is not None
|
|
assert retrieved.name == context.tracked_entity.name
|
|
|
|
|
|
@given('I have tracked a file entity named "{name}"')
|
|
def step_have_tracked_file_entity(context: Any, name: str) -> None:
|
|
"""Track a file entity."""
|
|
context.tracked_entity = context.memory_service.track_entity(name, EntityType.FILE)
|
|
context.original_first_seen = context.tracked_entity.first_seen
|
|
|
|
|
|
@when("I track the same entity again")
|
|
def step_track_same_entity(context: Any) -> None:
|
|
"""Track the same entity again to update it."""
|
|
time.sleep(0.01) # Small delay to ensure time difference
|
|
context.tracked_entity = context.memory_service.track_entity(
|
|
context.tracked_entity.name, context.tracked_entity.entity_type
|
|
)
|
|
|
|
|
|
@then("the entity last_seen should be updated")
|
|
def step_entity_last_seen_updated(context: Any) -> None:
|
|
"""Verify last_seen was updated."""
|
|
# last_seen should be >= first_seen
|
|
assert context.tracked_entity.last_seen >= context.original_first_seen
|
|
|
|
|
|
@given("a memory service with multiple entities tracked")
|
|
def step_memory_service_with_multiple_entities(context: Any) -> None:
|
|
"""Create memory service and track multiple entities."""
|
|
context.memory_service = MemoryService(session_id="multi-entity-session")
|
|
context.memory_service.track_entity("project-alpha", EntityType.PROJECT)
|
|
context.memory_service.track_entity("project-beta", EntityType.PROJECT)
|
|
context.memory_service.track_entity("feature-plan", EntityType.PLAN)
|
|
context.memory_service.track_entity("main.py", EntityType.FILE)
|
|
context.memory_service.track_entity("utils.py", EntityType.FILE)
|
|
|
|
|
|
@when('I get all entities of type "{entity_type}"')
|
|
def step_get_entities_by_type(context: Any, entity_type: str) -> None:
|
|
"""Get entities filtered by type."""
|
|
context.filtered_entities = context.memory_service.get_entities(entity_type)
|
|
|
|
|
|
@then("I should receive only project entities")
|
|
def step_receive_only_project_entities(context: Any) -> None:
|
|
"""Verify only project entities returned."""
|
|
assert len(context.filtered_entities) == 2
|
|
for entity in context.filtered_entities:
|
|
assert entity.entity_type == EntityType.PROJECT
|
|
|
|
|
|
@then("I should not receive entities of other types")
|
|
def step_not_receive_other_types(context: Any) -> None:
|
|
"""Verify no other entity types returned."""
|
|
for entity in context.filtered_entities:
|
|
assert entity.entity_type != EntityType.PLAN
|
|
assert entity.entity_type != EntityType.FILE
|
|
|
|
|
|
@when("I get all tracked entities")
|
|
def step_get_all_entities(context: Any) -> None:
|
|
"""Get all entities."""
|
|
context.all_entities = context.memory_service.get_entities()
|
|
|
|
|
|
@then("I should receive all entities regardless of type")
|
|
def step_receive_all_entities(context: Any) -> None:
|
|
"""Verify all entities returned."""
|
|
assert len(context.all_entities) == 5
|
|
types = {e.entity_type for e in context.all_entities}
|
|
assert EntityType.PROJECT in types
|
|
assert EntityType.PLAN in types
|
|
assert EntityType.FILE in types
|
|
|
|
|
|
@given("a memory service with multiple entities tracked over time")
|
|
def step_memory_service_entities_over_time(context: Any) -> None:
|
|
"""Track entities at different times."""
|
|
context.memory_service = MemoryService(session_id="time-entity-session")
|
|
|
|
# Track entities with small delays to ensure different timestamps
|
|
context.memory_service.track_entity("oldest", EntityType.PROJECT)
|
|
time.sleep(0.01)
|
|
context.memory_service.track_entity("middle", EntityType.PLAN)
|
|
time.sleep(0.01)
|
|
context.memory_service.track_entity("newest", EntityType.FILE)
|
|
|
|
|
|
@when("I get the {count:d} most recent entities")
|
|
def step_get_recent_entities(context: Any, count: int) -> None:
|
|
"""Get recent entities."""
|
|
context.recent_entities = context.memory_service.get_recent_entities(count)
|
|
|
|
|
|
@then("I should receive entities sorted by last access time")
|
|
def step_entities_sorted_by_time(context: Any) -> None:
|
|
"""Verify sorting by access time."""
|
|
# First should be most recent
|
|
assert context.recent_entities[0].name == "newest"
|
|
|
|
|
|
@then("I should receive at most {count:d} entities")
|
|
def step_receive_at_most_count(context: Any, count: int) -> None:
|
|
"""Verify count limit."""
|
|
assert len(context.recent_entities) <= count
|
|
|
|
|
|
@when('I search for entities with "{query}" in the name')
|
|
def step_search_entities(context: Any, query: str) -> None:
|
|
"""Search entities by name."""
|
|
context.search_results = context.memory_service.search_entities(query)
|
|
|
|
|
|
@then('I should receive entities whose names contain "{query}"')
|
|
def step_entities_contain_query(context: Any, query: str) -> None:
|
|
"""Verify search results contain query."""
|
|
for entity in context.search_results:
|
|
assert query.lower() in entity.name.lower()
|
|
|
|
|
|
@then("the search should be case-insensitive")
|
|
def step_search_case_insensitive(context: Any) -> None:
|
|
"""Verify case-insensitive search."""
|
|
# Search was done; if we got results, it worked
|
|
assert context.search_results is not None
|
|
|
|
|
|
@given('I have tracked a context entity named "{name}"')
|
|
def step_have_tracked_context_entity(context: Any, name: str) -> None:
|
|
"""Track a context entity."""
|
|
context.tracked_entity = context.memory_service.track_entity(
|
|
name, EntityType.CONTEXT
|
|
)
|
|
|
|
|
|
@when("I remove the entity")
|
|
def step_remove_entity(context: Any) -> None:
|
|
"""Remove the tracked entity."""
|
|
context.remove_result = context.memory_service.remove_entity(
|
|
context.tracked_entity.name, context.tracked_entity.entity_type
|
|
)
|
|
|
|
|
|
@then("the entity should no longer exist")
|
|
def step_entity_no_longer_exists(context: Any) -> None:
|
|
"""Verify entity was removed."""
|
|
assert context.remove_result is True
|
|
|
|
|
|
@then("attempting to get the removed entity should return None")
|
|
def step_get_removed_entity_returns_none(context: Any) -> None:
|
|
"""Verify removed entity returns None."""
|
|
retrieved = context.memory_service.get_entity(
|
|
context.tracked_entity.name, context.tracked_entity.entity_type
|
|
)
|
|
assert retrieved is None
|
|
|
|
|
|
@when("I clear all entities")
|
|
def step_clear_all_entities(context: Any) -> None:
|
|
"""Clear all entities."""
|
|
context.cleared_count = context.memory_service.clear_entities()
|
|
|
|
|
|
@then("the entity store should be empty")
|
|
def step_entity_store_empty(context: Any) -> None:
|
|
"""Verify entity store is empty."""
|
|
assert context.cleared_count > 0
|
|
|
|
|
|
@then("get_entities should return an empty list")
|
|
def step_get_entities_empty(context: Any) -> None:
|
|
"""Verify empty entity list."""
|
|
entities = context.memory_service.get_entities()
|
|
assert len(entities) == 0
|
|
|
|
|
|
@when("I clear only project entities")
|
|
def step_clear_only_project_entities(context: Any) -> None:
|
|
"""Clear only project entities."""
|
|
context.cleared_count = context.memory_service.clear_entities(EntityType.PROJECT)
|
|
|
|
|
|
@then("project entities should be removed")
|
|
def step_project_entities_removed(context: Any) -> None:
|
|
"""Verify project entities removed."""
|
|
projects = context.memory_service.get_entities(EntityType.PROJECT)
|
|
assert len(projects) == 0
|
|
|
|
|
|
@then("entities of other types should remain")
|
|
def step_other_entities_remain(context: Any) -> None:
|
|
"""Verify other entities remain."""
|
|
all_entities = context.memory_service.get_entities()
|
|
assert len(all_entities) > 0
|
|
# Should have plans and files
|
|
types = {e.entity_type for e in all_entities}
|
|
assert EntityType.PLAN in types or EntityType.FILE in types
|
|
|
|
|
|
@when("I get the entity summary")
|
|
def step_get_entity_summary(context: Any) -> None:
|
|
"""Get entity summary."""
|
|
context.entity_summary = context.memory_service.get_entity_summary()
|
|
|
|
|
|
@then("the summary should include total entity count")
|
|
def step_summary_has_total_count(context: Any) -> None:
|
|
"""Verify total count in summary."""
|
|
assert "total_entities" in context.entity_summary
|
|
assert context.entity_summary["total_entities"] >= 0
|
|
|
|
|
|
@then("the summary should include counts by type")
|
|
def step_summary_has_counts_by_type(context: Any) -> None:
|
|
"""Verify counts by type in summary."""
|
|
assert "counts_by_type" in context.entity_summary
|
|
assert isinstance(context.entity_summary["counts_by_type"], dict)
|
|
|
|
|
|
@then("the summary should include recent entities")
|
|
def step_summary_has_recent_entities(context: Any) -> None:
|
|
"""Verify recent entities in summary."""
|
|
assert "recent_entities" in context.entity_summary
|
|
assert isinstance(context.entity_summary["recent_entities"], list)
|
|
|
|
|
|
@when('I track an entity with type specified as string "{type_str}"')
|
|
def step_track_entity_with_string_type(context: Any, type_str: str) -> None:
|
|
"""Track entity using string type."""
|
|
context.tracked_entity = context.memory_service.track_entity(
|
|
"string-type-entity", type_str
|
|
)
|
|
|
|
|
|
@then("the entity should be tracked with EntityType.CUSTOM")
|
|
def step_entity_tracked_as_custom(context: Any) -> None:
|
|
"""Verify entity has CUSTOM type."""
|
|
assert context.tracked_entity.entity_type == EntityType.CUSTOM
|
|
|
|
|
|
@then("the entity should be retrievable using either string or enum type")
|
|
def step_entity_retrievable_with_either_type(context: Any) -> None:
|
|
"""Verify retrieval works with string or enum."""
|
|
# Using string
|
|
by_string = context.memory_service.get_entity("string-type-entity", "custom")
|
|
assert by_string is not None
|
|
|
|
# Using enum
|
|
by_enum = context.memory_service.get_entity("string-type-entity", EntityType.CUSTOM)
|
|
assert by_enum is not None
|
|
|
|
|
|
@when("I access the entity_store property")
|
|
def step_access_entity_store(context: Any) -> None:
|
|
"""Access entity store property."""
|
|
context.entity_store = context.memory_service.entity_store
|
|
|
|
|
|
@then("I should receive the underlying EntityStore instance")
|
|
def step_receive_entity_store_instance(context: Any) -> None:
|
|
"""Verify EntityStore instance returned."""
|
|
assert context.entity_store is not None
|
|
assert isinstance(context.entity_store, EntityStore)
|
|
|
|
|
|
@then("I can track entities through it directly")
|
|
def step_track_through_entity_store(context: Any) -> None:
|
|
"""Track entity directly through store."""
|
|
entity = context.entity_store.track("direct-entity", EntityType.FILE)
|
|
assert entity is not None
|
|
assert entity.name == "direct-entity"
|
|
|
|
|
|
# ==================== Coverage Enhancement Steps ====================
|
|
|
|
|
|
@given("an entity dictionary with valid data")
|
|
def step_entity_dictionary_with_valid_data(context: Any) -> None:
|
|
"""Create a dictionary representing a serialized entity."""
|
|
context.entity_dict = {
|
|
"name": "test-entity",
|
|
"entity_type": "project",
|
|
"metadata": {"key": "value", "priority": "high"},
|
|
"first_seen": "2024-01-15T10:30:00+00:00",
|
|
"last_seen": "2024-01-15T12:45:00+00:00",
|
|
"mention_count": 5,
|
|
}
|
|
|
|
|
|
@when("I create an Entity from the dictionary using from_dict")
|
|
def step_create_entity_from_dict(context: Any) -> None:
|
|
"""Deserialize entity from dictionary using from_dict classmethod."""
|
|
context.deserialized_entity = Entity.from_dict(context.entity_dict)
|
|
|
|
|
|
@then("the Entity should have the correct name and type")
|
|
def step_entity_has_correct_name_and_type(context: Any) -> None:
|
|
"""Verify entity name and type match."""
|
|
assert context.deserialized_entity.name == "test-entity"
|
|
assert context.deserialized_entity.entity_type == EntityType.PROJECT
|
|
|
|
|
|
@then("the Entity should have the correct metadata")
|
|
def step_entity_has_correct_metadata(context: Any) -> None:
|
|
"""Verify entity metadata matches."""
|
|
assert context.deserialized_entity.metadata == {"key": "value", "priority": "high"}
|
|
|
|
|
|
@then("the Entity should have the correct timestamps")
|
|
def step_entity_has_correct_timestamps(context: Any) -> None:
|
|
"""Verify entity timestamps are parsed correctly."""
|
|
from datetime import datetime
|
|
|
|
expected_first_seen = datetime(2024, 1, 15, 10, 30, 0, tzinfo=UTC)
|
|
expected_last_seen = datetime(2024, 1, 15, 12, 45, 0, tzinfo=UTC)
|
|
assert context.deserialized_entity.first_seen == expected_first_seen
|
|
assert context.deserialized_entity.last_seen == expected_last_seen
|
|
|
|
|
|
@then("the Entity should have the correct mention count")
|
|
def step_entity_has_correct_mention_count(context: Any) -> None:
|
|
"""Verify entity mention count matches."""
|
|
assert context.deserialized_entity.mention_count == 5
|
|
|
|
|
|
@given('I have tracked a project entity named "{name}" with initial metadata')
|
|
def step_track_project_with_initial_metadata(context: Any, name: str) -> None:
|
|
"""Track a project entity with initial metadata."""
|
|
context.initial_metadata = {"status": "draft", "created_by": "test"}
|
|
context.tracked_entity = context.memory_service.track_entity(
|
|
name, EntityType.PROJECT, context.initial_metadata
|
|
)
|
|
context.entity_name = name
|
|
|
|
|
|
@when("I track the same project entity again with additional metadata")
|
|
def step_track_same_entity_with_additional_metadata(context: Any) -> None:
|
|
"""Track the same entity with additional metadata to trigger update."""
|
|
time.sleep(0.01) # Small delay to ensure time difference
|
|
additional_metadata = {"status": "active", "priority": "high"}
|
|
context.tracked_entity = context.memory_service.track_entity(
|
|
context.entity_name, EntityType.PROJECT, additional_metadata
|
|
)
|
|
|
|
|
|
@then("the entity metadata should contain both initial and additional values")
|
|
def step_entity_metadata_contains_both(context: Any) -> None:
|
|
"""Verify metadata was merged correctly."""
|
|
# Original key should remain
|
|
assert context.tracked_entity.metadata.get("created_by") == "test"
|
|
# Status should be updated to new value
|
|
assert context.tracked_entity.metadata.get("status") == "active"
|
|
# New key should be added
|
|
assert context.tracked_entity.metadata.get("priority") == "high"
|
|
|
|
|
|
@when("I attempt to remove an entity that does not exist")
|
|
def step_attempt_remove_nonexistent_entity(context: Any) -> None:
|
|
"""Try to remove an entity that was never tracked."""
|
|
context.remove_result = context.memory_service.remove_entity(
|
|
"nonexistent-entity", EntityType.PROJECT
|
|
)
|
|
|
|
|
|
@then("the remove operation should return False")
|
|
def step_remove_returns_false(context: Any) -> None:
|
|
"""Verify remove returned False for missing entity."""
|
|
assert context.remove_result is False
|
|
|
|
|
|
@given("an entity store with a mock connection string")
|
|
def step_entity_store_with_connection_string(context: Any) -> None:
|
|
"""Create an entity store with a connection string to enable persistence."""
|
|
# Using a mock connection string to trigger persistence logic
|
|
context.entity_store = EntityStore(
|
|
session_id="persist-test",
|
|
connection_string="sqlite:///mock_persist.db",
|
|
)
|
|
|
|
|
|
@when("I track an entity to trigger dirty state")
|
|
def step_track_entity_trigger_dirty(context: Any) -> None:
|
|
"""Track an entity which sets dirty flag and triggers persistence."""
|
|
context.tracked_entity = context.entity_store.track(
|
|
"persist-entity", EntityType.PLAN
|
|
)
|
|
|
|
|
|
@then("the persist callback should mark the store as not dirty")
|
|
def step_persist_marks_not_dirty(context: Any) -> None:
|
|
"""Verify the store is not dirty after persistence."""
|
|
# After tracking, persist_if_needed is called, which clears dirty flag
|
|
assert context.entity_store._dirty is False
|
|
|
|
|
|
@then("subsequent reads should find the entity")
|
|
def step_subsequent_reads_find_entity(context: Any) -> None:
|
|
"""Verify the entity is readable after persistence."""
|
|
entity = context.entity_store.get("persist-entity", EntityType.PLAN)
|
|
assert entity is not None
|
|
assert entity.name == "persist-entity"
|
|
|
|
|
|
@given("a memory service with no messages")
|
|
def step_memory_service_no_messages(context: Any) -> None:
|
|
"""Create a memory service with empty message history."""
|
|
context.memory_service = MemoryService(session_id="empty-session")
|
|
|
|
|
|
@when("I request the conversation summary")
|
|
def step_request_conversation_summary(context: Any) -> None:
|
|
"""Get the summary from the memory service."""
|
|
context.summary = context.memory_service.get_summary()
|
|
|
|
|
|
@then("the summary should be an empty string")
|
|
def step_summary_is_empty_string(context: Any) -> None:
|
|
"""Verify the summary is an empty string."""
|
|
assert context.summary == ""
|