test: use scoped_handlers while listener Memory events

This commit is contained in:
Lucas Gomide
2025-07-01 23:07:56 -03:00
parent 15c5654a70
commit 174215980b
4 changed files with 119 additions and 124 deletions

View File

@@ -2499,58 +2499,59 @@ def test_using_contextual_memory():
def test_memory_events_are_emitted(): def test_memory_events_are_emitted():
events = defaultdict(list) events = defaultdict(list)
@crewai_event_bus.on(MemorySaveStartedEvent) with crewai_event_bus.scoped_handlers():
def handle_memory_save_started(source, event): @crewai_event_bus.on(MemorySaveStartedEvent)
events["MemorySaveStartedEvent"].append(event) def handle_memory_save_started(source, event):
events["MemorySaveStartedEvent"].append(event)
@crewai_event_bus.on(MemorySaveCompletedEvent) @crewai_event_bus.on(MemorySaveCompletedEvent)
def handle_memory_save_completed(source, event): def handle_memory_save_completed(source, event):
events["MemorySaveCompletedEvent"].append(event) events["MemorySaveCompletedEvent"].append(event)
@crewai_event_bus.on(MemorySaveFailedEvent) @crewai_event_bus.on(MemorySaveFailedEvent)
def handle_memory_save_failed(source, event): def handle_memory_save_failed(source, event):
events["MemorySaveFailedEvent"].append(event) events["MemorySaveFailedEvent"].append(event)
@crewai_event_bus.on(MemoryQueryStartedEvent) @crewai_event_bus.on(MemoryQueryStartedEvent)
def handle_memory_query_started(source, event): def handle_memory_query_started(source, event):
events["MemoryQueryStartedEvent"].append(event) events["MemoryQueryStartedEvent"].append(event)
@crewai_event_bus.on(MemoryQueryCompletedEvent) @crewai_event_bus.on(MemoryQueryCompletedEvent)
def handle_memory_query_completed(source, event): def handle_memory_query_completed(source, event):
events["MemoryQueryCompletedEvent"].append(event) events["MemoryQueryCompletedEvent"].append(event)
@crewai_event_bus.on(MemoryQueryFailedEvent) @crewai_event_bus.on(MemoryQueryFailedEvent)
def handle_memory_query_failed(source, event): def handle_memory_query_failed(source, event):
events["MemoryQueryFailedEvent"].append(event) events["MemoryQueryFailedEvent"].append(event)
@crewai_event_bus.on(MemoryRetrievalStartedEvent) @crewai_event_bus.on(MemoryRetrievalStartedEvent)
def handle_memory_retrieval_started(source, event): def handle_memory_retrieval_started(source, event):
events["MemoryRetrievalStartedEvent"].append(event) events["MemoryRetrievalStartedEvent"].append(event)
@crewai_event_bus.on(MemoryRetrievalCompletedEvent) @crewai_event_bus.on(MemoryRetrievalCompletedEvent)
def handle_memory_retrieval_completed(source, event): def handle_memory_retrieval_completed(source, event):
events["MemoryRetrievalCompletedEvent"].append(event) events["MemoryRetrievalCompletedEvent"].append(event)
math_researcher = Agent( math_researcher = Agent(
role="Researcher", role="Researcher",
goal="You research about math.", goal="You research about math.",
backstory="You're an expert in research and you love to learn new things.", backstory="You're an expert in research and you love to learn new things.",
allow_delegation=False, allow_delegation=False,
) )
task1 = Task( task1 = Task(
description="Research a topic to teach a kid aged 6 about math.", description="Research a topic to teach a kid aged 6 about math.",
expected_output="A topic, explanation, angle, and examples.", expected_output="A topic, explanation, angle, and examples.",
agent=math_researcher, agent=math_researcher,
) )
crew = Crew( crew = Crew(
agents=[math_researcher], agents=[math_researcher],
tasks=[task1], tasks=[task1],
memory=True, memory=True,
) )
crew.kickoff() crew.kickoff()
assert len(events["MemorySaveStartedEvent"]) == 6 assert len(events["MemorySaveStartedEvent"]) == 6
assert len(events["MemorySaveCompletedEvent"]) == 6 assert len(events["MemorySaveCompletedEvent"]) == 6

View File

@@ -238,20 +238,20 @@ def test_external_memory_search_events(custom_storage, external_memory_with_mock
events = defaultdict(list) events = defaultdict(list)
external_memory_with_mocked_config.storage = custom_storage external_memory_with_mocked_config.storage = custom_storage
with crewai_event_bus.scoped_handlers():
@crewai_event_bus.on(MemoryQueryStartedEvent)
def on_search_started(source, event):
events["MemoryQueryStartedEvent"].append(event)
@crewai_event_bus.on(MemoryQueryStartedEvent) @crewai_event_bus.on(MemoryQueryCompletedEvent)
def on_search_started(source, event): def on_search_completed(source, event):
events["MemoryQueryStartedEvent"].append(event) events["MemoryQueryCompletedEvent"].append(event)
@crewai_event_bus.on(MemoryQueryCompletedEvent) external_memory_with_mocked_config.search(
def on_search_completed(source, event): query="test value",
events["MemoryQueryCompletedEvent"].append(event) limit=3,
score_threshold=0.35,
external_memory_with_mocked_config.search( )
query="test value",
limit=3,
score_threshold=0.35,
)
assert len(events["MemoryQueryStartedEvent"]) == 1 assert len(events["MemoryQueryStartedEvent"]) == 1
assert len(events["MemoryQueryCompletedEvent"]) == 1 assert len(events["MemoryQueryCompletedEvent"]) == 1
@@ -281,7 +281,6 @@ def test_external_memory_search_events(custom_storage, external_memory_with_mock
'query_time_ms': ANY 'query_time_ms': ANY
} }
crewai_event_bus._handlers.clear()
def test_external_memory_save_events(custom_storage, external_memory_with_mocked_config): def test_external_memory_save_events(custom_storage, external_memory_with_mocked_config):
@@ -289,19 +288,20 @@ def test_external_memory_save_events(custom_storage, external_memory_with_mocked
external_memory_with_mocked_config.storage = custom_storage external_memory_with_mocked_config.storage = custom_storage
@crewai_event_bus.on(MemorySaveStartedEvent) with crewai_event_bus.scoped_handlers():
def on_save_started(source, event): @crewai_event_bus.on(MemorySaveStartedEvent)
events["MemorySaveStartedEvent"].append(event) def on_save_started(source, event):
events["MemorySaveStartedEvent"].append(event)
@crewai_event_bus.on(MemorySaveCompletedEvent) @crewai_event_bus.on(MemorySaveCompletedEvent)
def on_save_completed(source, event): def on_save_completed(source, event):
events["MemorySaveCompletedEvent"].append(event) events["MemorySaveCompletedEvent"].append(event)
external_memory_with_mocked_config.save( external_memory_with_mocked_config.save(
value="saving value", value="saving value",
metadata={"task": "test_task"}, metadata={"task": "test_task"},
agent="test_agent", agent="test_agent",
) )
assert len(events["MemorySaveStartedEvent"]) == 1 assert len(events["MemorySaveStartedEvent"]) == 1
assert len(events["MemorySaveCompletedEvent"]) == 1 assert len(events["MemorySaveCompletedEvent"]) == 1
@@ -329,5 +329,3 @@ def test_external_memory_save_events(custom_storage, external_memory_with_mocked
'agent_role': "test_agent", 'agent_role': "test_agent",
'save_time_ms': ANY 'save_time_ms': ANY
} }
crewai_event_bus._handlers.clear()

View File

@@ -20,23 +20,24 @@ def long_term_memory():
def test_long_term_memory_save_events(long_term_memory): def test_long_term_memory_save_events(long_term_memory):
events = defaultdict(list) events = defaultdict(list)
@crewai_event_bus.on(MemorySaveStartedEvent) with crewai_event_bus.scoped_handlers():
def on_save_started(source, event): @crewai_event_bus.on(MemorySaveStartedEvent)
events["MemorySaveStartedEvent"].append(event) def on_save_started(source, event):
events["MemorySaveStartedEvent"].append(event)
@crewai_event_bus.on(MemorySaveCompletedEvent) @crewai_event_bus.on(MemorySaveCompletedEvent)
def on_save_completed(source, event): def on_save_completed(source, event):
events["MemorySaveCompletedEvent"].append(event) events["MemorySaveCompletedEvent"].append(event)
memory = LongTermMemoryItem( memory = LongTermMemoryItem(
agent="test_agent", agent="test_agent",
task="test_task", task="test_task",
expected_output="test_output", expected_output="test_output",
datetime="test_datetime", datetime="test_datetime",
quality=0.5, quality=0.5,
metadata={"task": "test_task", "quality": 0.5}, metadata={"task": "test_task", "quality": 0.5},
) )
long_term_memory.save(memory) long_term_memory.save(memory)
assert len(events["MemorySaveStartedEvent"]) == 1 assert len(events["MemorySaveStartedEvent"]) == 1
assert len(events["MemorySaveCompletedEvent"]) == 1 assert len(events["MemorySaveCompletedEvent"]) == 1
@@ -64,26 +65,25 @@ def test_long_term_memory_save_events(long_term_memory):
"save_time_ms": ANY, "save_time_ms": ANY,
} }
crewai_event_bus._handlers.clear()
def test_long_term_memory_search_events(long_term_memory): def test_long_term_memory_search_events(long_term_memory):
events = defaultdict(list) events = defaultdict(list)
@crewai_event_bus.on(MemoryQueryStartedEvent) with crewai_event_bus.scoped_handlers():
def on_search_started(source, event): @crewai_event_bus.on(MemoryQueryStartedEvent)
events["MemoryQueryStartedEvent"].append(event) def on_search_started(source, event):
events["MemoryQueryStartedEvent"].append(event)
@crewai_event_bus.on(MemoryQueryCompletedEvent) @crewai_event_bus.on(MemoryQueryCompletedEvent)
def on_search_completed(source, event): def on_search_completed(source, event):
events["MemoryQueryCompletedEvent"].append(event) events["MemoryQueryCompletedEvent"].append(event)
test_query = "test query" test_query = "test query"
long_term_memory.search( long_term_memory.search(
test_query, test_query,
latest_n=5 latest_n=5
) )
assert len(events["MemoryQueryStartedEvent"]) == 1 assert len(events["MemoryQueryStartedEvent"]) == 1
assert len(events["MemoryQueryCompletedEvent"]) == 1 assert len(events["MemoryQueryCompletedEvent"]) == 1
@@ -113,7 +113,6 @@ def test_long_term_memory_search_events(long_term_memory):
'query_time_ms': ANY 'query_time_ms': ANY
} }
crewai_event_bus._handlers.clear()
def test_save_and_search(long_term_memory): def test_save_and_search(long_term_memory):
memory = LongTermMemoryItem( memory = LongTermMemoryItem(

View File

@@ -38,20 +38,21 @@ def short_term_memory():
def test_short_term_memory_search_events(short_term_memory): def test_short_term_memory_search_events(short_term_memory):
events = defaultdict(list) events = defaultdict(list)
@crewai_event_bus.on(MemoryQueryStartedEvent) with crewai_event_bus.scoped_handlers():
def on_search_started(source, event): @crewai_event_bus.on(MemoryQueryStartedEvent)
events["MemoryQueryStartedEvent"].append(event) def on_search_started(source, event):
events["MemoryQueryStartedEvent"].append(event)
@crewai_event_bus.on(MemoryQueryCompletedEvent) @crewai_event_bus.on(MemoryQueryCompletedEvent)
def on_search_completed(source, event): def on_search_completed(source, event):
events["MemoryQueryCompletedEvent"].append(event) events["MemoryQueryCompletedEvent"].append(event)
# Call the save method # Call the save method
short_term_memory.search( short_term_memory.search(
query="test value", query="test value",
limit=3, limit=3,
score_threshold=0.35, score_threshold=0.35,
) )
assert len(events["MemoryQueryStartedEvent"]) == 1 assert len(events["MemoryQueryStartedEvent"]) == 1
assert len(events["MemoryQueryCompletedEvent"]) == 1 assert len(events["MemoryQueryCompletedEvent"]) == 1
@@ -81,25 +82,23 @@ def test_short_term_memory_search_events(short_term_memory):
'query_time_ms': ANY 'query_time_ms': ANY
} }
crewai_event_bus._handlers.clear()
def test_short_term_memory_save_events(short_term_memory): def test_short_term_memory_save_events(short_term_memory):
events = defaultdict(list) events = defaultdict(list)
with crewai_event_bus.scoped_handlers():
@crewai_event_bus.on(MemorySaveStartedEvent)
def on_save_started(source, event):
events["MemorySaveStartedEvent"].append(event)
@crewai_event_bus.on(MemorySaveStartedEvent) @crewai_event_bus.on(MemorySaveCompletedEvent)
def on_save_started(source, event): def on_save_completed(source, event):
events["MemorySaveStartedEvent"].append(event) events["MemorySaveCompletedEvent"].append(event)
@crewai_event_bus.on(MemorySaveCompletedEvent) short_term_memory.save(
def on_save_completed(source, event): value="test value",
events["MemorySaveCompletedEvent"].append(event) metadata={"task": "test_task"},
agent="test_agent",
short_term_memory.save( )
value="test value",
metadata={"task": "test_task"},
agent="test_agent",
)
assert len(events["MemorySaveStartedEvent"]) == 1 assert len(events["MemorySaveStartedEvent"]) == 1
assert len(events["MemorySaveCompletedEvent"]) == 1 assert len(events["MemorySaveCompletedEvent"]) == 1
@@ -128,8 +127,6 @@ def test_short_term_memory_save_events(short_term_memory):
'save_time_ms': ANY 'save_time_ms': ANY
} }
crewai_event_bus._handlers.clear()
def test_save_and_search(short_term_memory): def test_save_and_search(short_term_memory):
memory = ShortTermMemoryItem( memory = ShortTermMemoryItem(
data="""test value test value test value test value test value test value data="""test value test value test value test value test value test value