mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-03 13:18:29 +00:00
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
import pytest
|
|
|
|
chromadb_not_installed = False
|
|
try:
|
|
import chromadb
|
|
except ImportError:
|
|
chromadb_not_installed = True
|
|
|
|
from crewai.memory.long_term.long_term_memory import LongTermMemory
|
|
from crewai.memory.long_term.long_term_memory_item import LongTermMemoryItem
|
|
|
|
|
|
@pytest.fixture
|
|
def long_term_memory():
|
|
"""Fixture to create a LongTermMemory instance"""
|
|
return LongTermMemory()
|
|
|
|
|
|
@pytest.mark.skipif(chromadb_not_installed, reason="ChromaDB is not installed")
|
|
def test_save_and_search(long_term_memory):
|
|
memory = LongTermMemoryItem(
|
|
agent="test_agent",
|
|
task="test_task",
|
|
expected_output="test_output",
|
|
datetime="test_datetime",
|
|
quality=0.5,
|
|
metadata={"task": "test_task", "quality": 0.5},
|
|
)
|
|
long_term_memory.save(memory)
|
|
find = long_term_memory.search("test_task", latest_n=5)[0]
|
|
assert find["score"] == 0.5
|
|
assert find["datetime"] == "test_datetime"
|
|
assert find["metadata"]["agent"] == "test_agent"
|
|
assert find["metadata"]["quality"] == 0.5
|
|
assert find["metadata"]["task"] == "test_task"
|
|
assert find["metadata"]["expected_output"] == "test_output"
|