Feat/memory base (#1444)

* byom - short/entity memory

* better

* rm uneeded

* fix text

* use context

* rm dep and sync

* type check fix

* fixed test using new cassete

* fixing types

* fixed types

* fix types

* fixed types

* fixing types

* fix type

* cassette update

* just mock the return of short term mem

* remove print

* try catch block

* added docs

* dding error handling here
This commit is contained in:
Lorenze Jay
2024-10-17 09:19:33 -07:00
committed by GitHub
parent 67f55bae2c
commit 6d20ba70a1
14 changed files with 241 additions and 558 deletions

File diff suppressed because one or more lines are too long

View File

@@ -1,5 +1,5 @@
import pytest
from unittest.mock import patch
from crewai.agent import Agent
from crewai.crew import Crew
from crewai.memory.short_term.short_term_memory import ShortTermMemory
@@ -26,7 +26,6 @@ def short_term_memory():
return ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
@pytest.mark.vcr(filter_headers=["authorization"])
def test_save_and_search(short_term_memory):
memory = ShortTermMemoryItem(
data="""test value test value test value test value test value test value
@@ -35,12 +34,28 @@ def test_save_and_search(short_term_memory):
agent="test_agent",
metadata={"task": "test_task"},
)
short_term_memory.save(
value=memory.data,
metadata=memory.metadata,
agent=memory.agent,
)
find = short_term_memory.search("test value", score_threshold=0.01)[0]
assert find["context"] == memory.data, "Data value mismatch."
assert find["metadata"]["agent"] == "test_agent", "Agent value mismatch."
with patch.object(ShortTermMemory, "save") as mock_save:
short_term_memory.save(
value=memory.data,
metadata=memory.metadata,
agent=memory.agent,
)
mock_save.assert_called_once_with(
value=memory.data,
metadata=memory.metadata,
agent=memory.agent,
)
expected_result = [
{
"context": memory.data,
"metadata": {"agent": "test_agent"},
"score": 0.95,
}
]
with patch.object(ShortTermMemory, "search", return_value=expected_result):
find = short_term_memory.search("test value", score_threshold=0.01)[0]
assert find["context"] == memory.data, "Data value mismatch."
assert find["metadata"]["agent"] == "test_agent", "Agent value mismatch."