mirror of
https://github.com/crewAIInc/crewAI.git
synced 2025-12-30 19:28:29 +00:00
131 lines
5.1 KiB
Python
131 lines
5.1 KiB
Python
import time
|
|
from datetime import datetime, timedelta
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from crewai.agent import Agent
|
|
from crewai.crew import Crew
|
|
from crewai.memory.short_term.short_term_memory import ShortTermMemory
|
|
from crewai.memory.short_term.short_term_memory_item import ShortTermMemoryItem
|
|
from crewai.memory.storage.rag_storage import RAGStorage
|
|
from crewai.task import Task
|
|
|
|
|
|
@pytest.fixture
|
|
def short_term_memory():
|
|
"""Fixture to create a ShortTermMemory instance"""
|
|
agent = Agent(
|
|
role="Tutor",
|
|
goal="Teach programming concepts",
|
|
backstory="You are a programming tutor helping students learn.",
|
|
tools=[],
|
|
verbose=True,
|
|
)
|
|
|
|
task = Task(
|
|
description="Explain programming concepts to students.",
|
|
expected_output="Clear explanations of programming concepts.",
|
|
agent=agent,
|
|
)
|
|
return ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
|
|
|
|
|
|
def test_memory_prioritizes_recent_topic(short_term_memory):
|
|
"""Test that memory retrieval prioritizes the most recent topic in a conversation."""
|
|
# First topic: Python variables
|
|
topic1_data = "Variables in Python are dynamically typed. You can assign any value to a variable without declaring its type."
|
|
topic1_timestamp = datetime.now() - timedelta(minutes=10) # Older memory
|
|
|
|
# Second topic: Python abstract classes
|
|
topic2_data = "Abstract classes in Python are created using the ABC module. They cannot be instantiated and are used as a blueprint for other classes."
|
|
topic2_timestamp = datetime.now() # More recent memory
|
|
|
|
# Mock search results to simulate what would be returned by RAGStorage
|
|
mock_results = [
|
|
{
|
|
"id": "2",
|
|
"metadata": {
|
|
"agent": "Tutor",
|
|
"topic": "python_abstract_classes",
|
|
"timestamp": topic2_timestamp.isoformat()
|
|
},
|
|
"context": topic2_data,
|
|
"score": 0.85, # Higher score due to recency boost
|
|
},
|
|
{
|
|
"id": "1",
|
|
"metadata": {
|
|
"agent": "Tutor",
|
|
"topic": "python_variables",
|
|
"timestamp": topic1_timestamp.isoformat()
|
|
},
|
|
"context": topic1_data,
|
|
"score": 0.75, # Lower score due to being older
|
|
}
|
|
]
|
|
|
|
# Mock the search method to return our predefined results
|
|
with patch.object(RAGStorage, 'search', return_value=mock_results):
|
|
# Query that could match both topics but should prioritize the more recent one
|
|
query = "Can you give me another example of that?"
|
|
|
|
# Search with recency consideration
|
|
results = short_term_memory.search(query)
|
|
|
|
# Verify that the most recent topic (abstract classes) is prioritized
|
|
assert len(results) > 0, "No search results returned"
|
|
|
|
# The first result should be about abstract classes (the more recent topic)
|
|
assert "abstract classes" in results[0]["context"].lower(), "Recent topic (abstract classes) not prioritized"
|
|
|
|
# If there are multiple results, check if the older topic is also returned but with lower priority
|
|
if len(results) > 1:
|
|
assert "variables" in results[1]["context"].lower(), "Older topic should be second"
|
|
|
|
# Verify that the scores reflect the recency prioritization
|
|
assert results[0]["score"] > results[1]["score"], "Recent topic should have higher score"
|
|
|
|
|
|
def test_future_timestamp_validation():
|
|
"""Test that ShortTermMemoryItem raises ValueError for future timestamps."""
|
|
# Setup agent and task for memory
|
|
agent = Agent(
|
|
role="Tutor",
|
|
goal="Teach programming concepts",
|
|
backstory="You are a programming tutor helping students learn.",
|
|
tools=[],
|
|
verbose=True,
|
|
)
|
|
|
|
task = Task(
|
|
description="Explain programming concepts to students.",
|
|
expected_output="Clear explanations of programming concepts.",
|
|
agent=agent,
|
|
)
|
|
|
|
# Create a future timestamp
|
|
future_timestamp = datetime.now() + timedelta(days=1)
|
|
|
|
# Test constructor validation
|
|
with pytest.raises(ValueError, match="Timestamp cannot be in the future"):
|
|
ShortTermMemoryItem(data="Test data", timestamp=future_timestamp)
|
|
|
|
# Test save method validation
|
|
memory = ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
|
|
|
|
# Create a memory item with a future timestamp
|
|
future_data = "Test data with future timestamp"
|
|
|
|
# We need to pass the data directly to the save method
|
|
# The save method will create a ShortTermMemoryItem internally
|
|
# and then we'll modify its timestamp before it's saved
|
|
|
|
# Mock datetime.now to return a fixed time
|
|
with patch('crewai.memory.short_term.short_term_memory_item.datetime') as mock_datetime:
|
|
# Set up the mock to return our future timestamp when now() is called
|
|
mock_datetime.now.return_value = future_timestamp
|
|
|
|
with pytest.raises(ValueError, match="Cannot save memory item with future timestamp"):
|
|
memory.save(value=future_data)
|