Minor Changes

This commit is contained in:
Vidit-Ostwal
2025-03-28 22:08:15 +05:30
parent 625748e462
commit 08a6a82071
5 changed files with 270 additions and 22 deletions

View File

@@ -164,7 +164,10 @@ crew = Crew(
[Mem0](https://mem0.ai/) is a self-improving memory layer for LLM applications, enabling personalized AI experiences.
To include user-specific memory you can get your API key [here](https://app.mem0.ai/dashboard/api-keys) and refer the [docs](https://docs.mem0.ai/platform/quickstart#4-1-create-memories) for adding user preferences.
### Using Mem0 API platform
To include user-specific memory you can get your API key [here](https://app.mem0.ai/dashboard/api-keys) and refer the [docs](https://docs.mem0.ai/platform/quickstart#4-1-create-memories) for adding user preferences. In this case `user_memory` is set to `MemoryClient` from mem0.
```python Code
@@ -175,18 +178,7 @@ from mem0 import MemoryClient
# Set environment variables for Mem0
os.environ["MEM0_API_KEY"] = "m0-xx"
# Step 1: Record preferences based on past conversation or user input
client = MemoryClient()
messages = [
{"role": "user", "content": "Hi there! I'm planning a vacation and could use some advice."},
{"role": "assistant", "content": "Hello! I'd be happy to help with your vacation planning. What kind of destination do you prefer?"},
{"role": "user", "content": "I am more of a beach person than a mountain person."},
{"role": "assistant", "content": "That's interesting. Do you like hotels or Airbnb?"},
{"role": "user", "content": "I like Airbnb more."},
]
client.add(messages, user_id="john")
# Step 2: Create a Crew with User Memory
# Step 1: Create a Crew with User Memory
crew = Crew(
agents=[...],
@@ -197,11 +189,12 @@ crew = Crew(
memory_config={
"provider": "mem0",
"config": {"user_id": "john"},
"user_memory : {}" #Set user_memory explicitly to a dictionary, we are working on this issue.
},
)
```
## Memory Configuration Options
#### Additional Memory Configuration Options
If you want to access a specific organization and project, you can set the `org_id` and `project_id` parameters in the memory configuration.
```python Code
@@ -215,10 +208,74 @@ crew = Crew(
memory_config={
"provider": "mem0",
"config": {"user_id": "john", "org_id": "my_org_id", "project_id": "my_project_id"},
"user_memory : {}" #Set user_memory explicitly to a dictionary, we are working on this issue.
},
)
```
### Using Local Mem0 memory
If you want to use local mem0 memory, with a custom configuration, you can set a parameter `local_mem0_config` in the config itself.
If both os environment key is set and local_mem0_config is given, the API platform takes higher priority over the local configuration.
Check [this](https://docs.mem0.ai/open-source/python-quickstart#run-mem0-locally) mem0 local configuration docs for more understanding.
In this case `user_memory` is set to `Memory` from mem0.
```python Code
from crewai import Crew
#local mem0 config
config = {
"vector_store": {
"provider": "qdrant",
"config": {
"host": "localhost",
"port": 6333
}
},
"llm": {
"provider": "openai",
"config": {
"api_key": "your-api-key",
"model": "gpt-4"
}
},
"embedder": {
"provider": "openai",
"config": {
"api_key": "your-api-key",
"model": "text-embedding-3-small"
}
},
"graph_store": {
"provider": "neo4j",
"config": {
"url": "neo4j+s://your-instance",
"username": "neo4j",
"password": "password"
}
},
"history_db_path": "/path/to/history.db",
"version": "v1.1",
"custom_fact_extraction_prompt": "Optional custom prompt for fact extraction for memory",
"custom_update_memory_prompt": "Optional custom prompt for update memory"
}
crew = Crew(
agents=[...],
tasks=[...],
verbose=True,
memory=True,
memory_config={
"provider": "mem0",
"config": {"user_id": "john", 'local_mem0_config': config},
"user_memory : {}" #Set user_memory explicitly to a dictionary, we are working on this issue.
},
)
```
## Additional Embedding Providers
### Using OpenAI embeddings (already default)

View File

@@ -290,24 +290,25 @@ class Crew(BaseModel):
else EntityMemory(crew=self, embedder_config=self.embedder)
)
if (
self.memory_config and "user_memory" in self.memory_config
self.memory_config and "user_memory" in self.memory_config and self.memory_config.get('provider') == 'mem0'
): # Check for user_memory in config
user_memory_config = self.memory_config["user_memory"]
if isinstance(
user_memory_config, UserMemory
): # Check if it is already an instance
self._user_memory = user_memory_config
elif isinstance(
user_memory_config, dict
): # Check if it's a configuration dict
self._user_memory = UserMemory(
crew=self, **user_memory_config
) # Initialize with config
crew=self
)
else:
raise TypeError(
"user_memory must be a UserMemory instance or a configuration dictionary"
"user_memory must be a configuration dictionary"
)
else:
self._logger.log(
"warning",
"User memory initialization failed. For setup instructions, please refer to the memory documentation: https://docs.crewai.com/concepts/memory#integrating-mem0-for-enhanced-user-memory",
color="yellow"
)
self._user_memory = None # No user memory if not in config
return self

View File

@@ -43,3 +43,11 @@ class UserMemory(Memory):
score_threshold=score_threshold,
)
return results
def reset(self) -> None:
try:
self.storage.reset()
except Exception as e:
raise Exception(
f"An error occurred while resetting the user memory: {e}"
)

View File

@@ -0,0 +1,68 @@
from unittest.mock import MagicMock, patch
import pytest
from mem0.memory.main import Memory
from crewai.memory.user.user_memory import UserMemory
from crewai.memory.user.user_memory_item import UserMemoryItem
class MockCrew:
def __init__(self, memory_config):
self.memory_config = memory_config
@pytest.fixture
def user_memory():
"""Fixture to create a UserMemory instance"""
crew = MockCrew(
memory_config={
"provider": "mem0",
"config": {"user_id": "john"},
"user_memory" : {}
}
)
user_memory = MagicMock(spec=UserMemory)
with patch.object(Memory,'__new__',return_value=user_memory):
user_memory_instance = UserMemory(crew=crew)
return user_memory_instance
def test_save_and_search(user_memory):
memory = UserMemoryItem(
data="""test value test value test value test value test value test value
test value test value test value test value test value test value
test value test value test value test value test value test value""",
user="test_user",
metadata={"task": "test_task"},
)
with patch.object(UserMemory, "save") as mock_save:
user_memory.save(
value=memory.data,
metadata=memory.metadata,
user=memory.user
)
mock_save.assert_called_once_with(
value=memory.data,
metadata=memory.metadata,
user=memory.user
)
expected_result = [
{
"context": memory.data,
"metadata": {"agent": "test_agent"},
"score": 0.95,
}
]
expected_result = ["mocked_result"]
# Use patch.object to mock UserMemory's search method
with patch.object(UserMemory, 'search', return_value=expected_result) as mock_search:
find = UserMemory.search("test value", score_threshold=0.01)[0]
mock_search.assert_called_once_with("test value", score_threshold=0.01)
assert find == expected_result[0]

View File

@@ -0,0 +1,114 @@
import os
from unittest.mock import MagicMock, patch
import pytest
from mem0.client.main import MemoryClient
from mem0.memory.main import Memory
from crewai.agent import Agent
from crewai.crew import Crew
from crewai.memory.storage.mem0_storage import Mem0Storage
from crewai.task import Task
# Define the class (if not already defined)
class MockCrew:
def __init__(self, memory_config):
self.memory_config = memory_config
@pytest.fixture
def mock_mem0_memory():
"""Fixture to create a mock Memory instance"""
mock_memory = MagicMock(spec=Memory)
return mock_memory
@pytest.fixture
def mem0_storage_with_mocked_config(mock_mem0_memory):
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
# Patch the Memory class to return our mock
with patch('mem0.memory.main.Memory.from_config', return_value=mock_mem0_memory):
config = {
"vector_store": {
"provider": "mock_vector_store",
"config": {
"host": "localhost",
"port": 6333
}
},
"llm": {
"provider": "mock_llm",
"config": {
"api_key": "mock-api-key",
"model": "mock-model"
}
},
"embedder": {
"provider": "mock_embedder",
"config": {
"api_key": "mock-api-key",
"model": "mock-model"
}
},
"graph_store": {
"provider": "mock_graph_store",
"config": {
"url": "mock-url",
"username": "mock-user",
"password": "mock-password"
}
},
"history_db_path": "/mock/path",
"version": "test-version",
"custom_fact_extraction_prompt": "mock prompt 1",
"custom_update_memory_prompt": "mock prompt 2"
}
# Instantiate the class with memory_config
crew = MockCrew(
memory_config={
"provider": "mem0",
"config": {"user_id": "test_user", "local_mem0_config": config},
}
)
mem0_storage = Mem0Storage(type="short_term", crew=crew)
return mem0_storage
def test_mem0_storage_initialization(mem0_storage_with_mocked_config, mock_mem0_memory):
"""Test that Mem0Storage initializes correctly with the mocked config"""
assert mem0_storage_with_mocked_config.memory_type == "short_term"
assert mem0_storage_with_mocked_config.memory is mock_mem0_memory
@pytest.fixture
def mock_mem0_memory_client():
"""Fixture to create a mock MemoryClient instance"""
mock_memory = MagicMock(spec=MemoryClient)
return mock_memory
@pytest.fixture
def mem0_storage_with_memory_client(mock_mem0_memory_client):
"""Fixture to create a Mem0Storage instance with mocked dependencies"""
# We need to patch the MemoryClient before it's instantiated
with patch.object(MemoryClient, '__new__', return_value=mock_mem0_memory_client):
crew = MockCrew(
memory_config={
"provider": "mem0",
"config": {"user_id": "test_user", "api_key": "ABCDEFGH", "org_id": "my_org_id", "project_id": "my_project_id"},
}
)
mem0_storage = Mem0Storage(type="short_term", crew=crew)
return mem0_storage
def test_mem0_storage_with_memory_client_initialization(mem0_storage_with_memory_client, mock_mem0_memory_client):
"""Test Mem0Storage initialization with MemoryClient"""
assert mem0_storage_with_memory_client.memory_type == "short_term"
assert mem0_storage_with_memory_client.memory is mock_mem0_memory_client