mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 16:48:30 +00:00
212 lines
7.3 KiB
Python
212 lines
7.3 KiB
Python
from typing import Any, Dict, List
|
|
|
|
import pytest
|
|
|
|
from crewai.agent import Agent
|
|
from crewai.crew import Crew
|
|
from crewai.memory.entity.entity_memory import EntityMemory
|
|
from crewai.memory.long_term.long_term_memory import LongTermMemory
|
|
from crewai.memory.short_term.short_term_memory import ShortTermMemory
|
|
from crewai.memory.storage.interface import Storage, SearchResult
|
|
from crewai.memory.user.user_memory import UserMemory
|
|
|
|
|
|
class CustomStorage(Storage):
|
|
"""Custom storage implementation for testing."""
|
|
|
|
def __init__(self):
|
|
self.data = []
|
|
|
|
def save(self, value: Any, metadata: Dict[str, Any]) -> None:
|
|
self.data.append({"value": value, "metadata": metadata})
|
|
|
|
def search(
|
|
self, query: str, limit: int = 3, score_threshold: float = 0.35
|
|
) -> List[SearchResult]:
|
|
return [{"context": item["value"], "metadata": item["metadata"], "score": 0.9} for item in self.data]
|
|
|
|
def reset(self) -> None:
|
|
self.data = []
|
|
|
|
|
|
def test_custom_storage_with_short_term_memory():
|
|
"""Test that custom storage works with short term memory."""
|
|
custom_storage = CustomStorage()
|
|
memory = ShortTermMemory(storage=custom_storage)
|
|
|
|
memory.save("test value", {"key": "value"})
|
|
results = memory.search("test")
|
|
|
|
assert len(results) > 0
|
|
assert results[0]["context"] == "test value"
|
|
assert results[0]["metadata"]["key"] == "value"
|
|
|
|
|
|
def test_custom_storage_with_long_term_memory():
|
|
"""Test that custom storage works with long term memory."""
|
|
custom_storage = CustomStorage()
|
|
memory = LongTermMemory(storage=custom_storage)
|
|
|
|
memory.save("test value", {"key": "value"})
|
|
results = memory.search("test")
|
|
|
|
assert len(results) > 0
|
|
assert results[0]["context"] == "test value"
|
|
assert results[0]["metadata"]["key"] == "value"
|
|
|
|
|
|
def test_custom_storage_with_entity_memory():
|
|
"""Test that custom storage works with entity memory."""
|
|
custom_storage = CustomStorage()
|
|
memory = EntityMemory(storage=custom_storage)
|
|
|
|
memory.save("test value", {"key": "value"})
|
|
results = memory.search("test")
|
|
|
|
assert len(results) > 0
|
|
assert results[0]["context"] == "test value"
|
|
assert results[0]["metadata"]["key"] == "value"
|
|
|
|
|
|
def test_custom_storage_with_user_memory():
|
|
"""Test that custom storage works with user memory."""
|
|
custom_storage = CustomStorage()
|
|
memory = UserMemory(storage=custom_storage)
|
|
|
|
memory.save("test value", {"key": "value"})
|
|
results = memory.search("test")
|
|
|
|
assert len(results) > 0
|
|
# UserMemory prepends "Remember the details about the user: " to the value
|
|
assert "test value" in results[0]["context"]
|
|
assert results[0]["metadata"]["key"] == "value"
|
|
|
|
|
|
def test_custom_storage_with_crew():
|
|
"""Test that custom storage works with crew."""
|
|
short_term_storage = CustomStorage()
|
|
long_term_storage = CustomStorage()
|
|
entity_storage = CustomStorage()
|
|
user_storage = CustomStorage()
|
|
|
|
# Create memory instances with custom storage
|
|
short_term_memory = ShortTermMemory(storage=short_term_storage)
|
|
long_term_memory = LongTermMemory(storage=long_term_storage)
|
|
entity_memory = EntityMemory(storage=entity_storage)
|
|
user_memory = UserMemory(storage=user_storage)
|
|
|
|
# Create a crew with custom memory instances
|
|
crew = Crew(
|
|
agents=[Agent(role="test", goal="test", backstory="test")],
|
|
memory=True,
|
|
short_term_memory=short_term_memory,
|
|
long_term_memory=long_term_memory,
|
|
entity_memory=entity_memory,
|
|
memory_config={"user_memory": user_memory},
|
|
)
|
|
|
|
# Test that the crew has the custom memory instances
|
|
assert crew._short_term_memory.storage == short_term_storage
|
|
assert crew._long_term_memory.storage == long_term_storage
|
|
assert crew._entity_memory.storage == entity_storage
|
|
assert crew._user_memory.storage == user_storage
|
|
|
|
|
|
def test_custom_storage_with_memory_config():
|
|
"""Test that custom storage works with memory_config."""
|
|
short_term_storage = CustomStorage()
|
|
long_term_memory = LongTermMemory(storage=CustomStorage())
|
|
entity_memory = EntityMemory(storage=CustomStorage())
|
|
user_memory = UserMemory(storage=CustomStorage())
|
|
|
|
# Create a crew with custom storage in memory_config
|
|
crew = Crew(
|
|
agents=[Agent(role="test", goal="test", backstory="test")],
|
|
memory=True,
|
|
short_term_memory=ShortTermMemory(storage=short_term_storage),
|
|
long_term_memory=long_term_memory,
|
|
entity_memory=entity_memory,
|
|
memory_config={
|
|
"user_memory": user_memory
|
|
},
|
|
)
|
|
|
|
# Test that the crew has the custom storage instances
|
|
assert crew._short_term_memory.storage == short_term_storage
|
|
assert crew._long_term_memory == long_term_memory
|
|
assert crew._entity_memory == entity_memory
|
|
assert crew._user_memory == user_memory
|
|
|
|
|
|
def test_custom_storage_error_handling():
|
|
"""Test error handling with custom storage."""
|
|
# Test exception propagation
|
|
class ErrorStorage(Storage):
|
|
"""Storage implementation that raises exceptions."""
|
|
def __init__(self):
|
|
self.data = []
|
|
|
|
def save(self, value: Any, metadata: Dict[str, Any]) -> None:
|
|
raise ValueError("Save error")
|
|
|
|
def search(
|
|
self, query: str, limit: int = 3, score_threshold: float = 0.35
|
|
) -> List[SearchResult]:
|
|
raise ValueError("Search error")
|
|
|
|
def reset(self) -> None:
|
|
raise ValueError("Reset error")
|
|
|
|
storage = ErrorStorage()
|
|
memory = ShortTermMemory(storage=storage)
|
|
|
|
with pytest.raises(ValueError, match="Save error"):
|
|
memory.save("test", {})
|
|
|
|
with pytest.raises(ValueError, match="Search error"):
|
|
memory.search("test")
|
|
|
|
with pytest.raises(Exception, match="An error occurred while resetting the short-term memory: Reset error"):
|
|
memory.reset()
|
|
|
|
|
|
def test_custom_storage_edge_cases():
|
|
"""Test edge cases with custom storage."""
|
|
class EdgeCaseStorage(Storage):
|
|
"""Storage implementation for testing edge cases."""
|
|
def __init__(self):
|
|
self.data = []
|
|
|
|
def save(self, value: Any, metadata: Dict[str, Any]) -> None:
|
|
self.data.append({"value": value, "metadata": metadata})
|
|
|
|
def search(
|
|
self, query: str, limit: int = 3, score_threshold: float = 0.35
|
|
) -> List[SearchResult]:
|
|
return [{"context": item["value"], "metadata": item["metadata"], "score": 0.5} for item in self.data]
|
|
|
|
def reset(self) -> None:
|
|
self.data = []
|
|
|
|
storage = EdgeCaseStorage()
|
|
memory = ShortTermMemory(storage=storage)
|
|
|
|
# Test empty query
|
|
memory.save("test value", {"key": "value"})
|
|
results = memory.search("")
|
|
assert len(results) > 0
|
|
|
|
# Test very large metadata
|
|
large_metadata = {"key" + str(i): "value" * 100 for i in range(100)}
|
|
memory.save("test value", large_metadata)
|
|
results = memory.search("test")
|
|
assert len(results) > 0
|
|
assert results[1]["metadata"] == large_metadata
|
|
|
|
# Test unicode and special characters
|
|
unicode_value = "测试值 with special chars: !@#$%^&*()"
|
|
memory.save(unicode_value, {"key": "value"})
|
|
results = memory.search("测试")
|
|
assert len(results) > 0
|
|
assert unicode_value in results[2]["context"]
|