klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Code formatting (#1986)

Browse Source
This commit is contained in:
Dev Khant
2024-10-29 11:32:07 +05:30
committed by GitHub
parent dca74a1ec0
commit 605558da9d
13 changed files with 119 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
cookbooks/helper/mem0_teachability.py
View File

@@ -12,7 +12,7 @@ from autogen.agentchat.contrib.capabilities.agent_capability import AgentCapabil
from autogen.agentchat.contrib.text_analyzer_agent import TextAnalyzerAgent from autogen.agentchat.contrib.text_analyzer_agent import TextAnalyzerAgent
from termcolor import colored from termcolor import colored
from mem0 import Memory from mem0 import Memory
from mem0.configs.base import MemoryConfig
class Mem0Teachability(AgentCapability): class Mem0Teachability(AgentCapability):
def __init__( def __init__(
@@ -60,7 +60,6 @@ class Mem0Teachability(AgentCapability):
return expanded_text return expanded_text
def _consider_memo_storage(self, comment: Union[Dict, str]): def _consider_memo_storage(self, comment: Union[Dict, str]):
memo_added = False
response = self._analyze( response = self._analyze(
comment, comment,
"Does any part of the TEXT ask the agent to perform a task or solve a problem? Answer with just one word, yes or no.", "Does any part of the TEXT ask the agent to perform a task or solve a problem? Answer with just one word, yes or no.",
@@ -85,8 +84,9 @@ class Mem0Teachability(AgentCapability):
if self.verbosity >= 1: if self.verbosity >= 1:
print(colored("\nREMEMBER THIS TASK-ADVICE PAIR", "light_yellow")) print(colored("\nREMEMBER THIS TASK-ADVICE PAIR", "light_yellow"))
self.memory.add([{"role": "user", "content": f"Task: {general_task}\nAdvice: {advice}"}], agent_id=self.agent_id) self.memory.add(
memo_added = True [{"role": "user", "content": f"Task: {general_task}\nAdvice: {advice}"}], agent_id=self.agent_id
)
response = self._analyze( response = self._analyze(
comment, comment,
@@ -105,8 +105,9 @@ class Mem0Teachability(AgentCapability):
if self.verbosity >= 1: if self.verbosity >= 1:
print(colored("\nREMEMBER THIS QUESTION-ANSWER PAIR", "light_yellow")) print(colored("\nREMEMBER THIS QUESTION-ANSWER PAIR", "light_yellow"))
self.memory.add([{"role": "user", "content": f"Question: {question}\nAnswer: {answer}"}], agent_id=self.agent_id) self.memory.add(
memo_added = True [{"role": "user", "content": f"Question: {question}\nAnswer: {answer}"}], agent_id=self.agent_id
)
def _consider_memo_retrieval(self, comment: Union[Dict, str]): def _consider_memo_retrieval(self, comment: Union[Dict, str]):
if self.verbosity >= 1: if self.verbosity >= 1:

46
mem0/client/main.py
View File

@@ -117,7 +117,9 @@ class MemoryClient:
payload = self._prepare_payload(messages, kwargs) payload = self._prepare_payload(messages, kwargs)
response = self.client.post("/v1/memories/", json=payload) response = self.client.post("/v1/memories/", json=payload)
response.raise_for_status() response.raise_for_status()
capture_client_event("client.add", self) if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event("client.add", self, {"keys": list(kwargs.keys())})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -135,7 +137,7 @@ class MemoryClient:
""" """
response = self.client.get(f"/v1/memories/{memory_id}/") response = self.client.get(f"/v1/memories/{memory_id}/")
response.raise_for_status() response.raise_for_status()
capture_client_event("client.get", self) capture_client_event("client.get", self, {"memory_id": memory_id})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -159,10 +161,12 @@ class MemoryClient:
elif version == "v2": elif version == "v2":
response = self.client.post(f"/{version}/memories/", json=params) response = self.client.post(f"/{version}/memories/", json=params)
response.raise_for_status() response.raise_for_status()
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event( capture_client_event(
"client.get_all", "client.get_all",
self, self,
{"filters": len(params), "limit": kwargs.get("limit", 100)}, {"api_version": version, "keys": list(kwargs.keys())},
) )
return response.json() return response.json()
@@ -186,7 +190,9 @@ class MemoryClient:
payload.update({k: v for k, v in kwargs.items() if v is not None}) payload.update({k: v for k, v in kwargs.items() if v is not None})
response = self.client.post(f"/{version}/memories/search/", json=payload) response = self.client.post(f"/{version}/memories/search/", json=payload)
response.raise_for_status() response.raise_for_status()
capture_client_event("client.search", self, {"limit": kwargs.get("limit", 100)}) if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event("client.search", self, {"api_version": version, "keys": list(kwargs.keys())})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -199,7 +205,7 @@ class MemoryClient:
Returns: Returns:
Dict[str, Any]: The response from the server. Dict[str, Any]: The response from the server.
""" """
capture_client_event("client.update", self) capture_client_event("client.update", self, {"memory_id": memory_id})
response = self.client.put(f"/v1/memories/{memory_id}/", json={"text": data}) response = self.client.put(f"/v1/memories/{memory_id}/", json={"text": data})
response.raise_for_status() response.raise_for_status()
return response.json() return response.json()
@@ -219,7 +225,7 @@ class MemoryClient:
""" """
response = self.client.delete(f"/v1/memories/{memory_id}/") response = self.client.delete(f"/v1/memories/{memory_id}/")
response.raise_for_status() response.raise_for_status()
capture_client_event("client.delete", self) capture_client_event("client.delete", self, {"memory_id": memory_id})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -239,7 +245,7 @@ class MemoryClient:
params = self._prepare_params(kwargs) params = self._prepare_params(kwargs)
response = self.client.delete("/v1/memories/", params=params) response = self.client.delete("/v1/memories/", params=params)
response.raise_for_status() response.raise_for_status()
capture_client_event("client.delete_all", self, {"params": len(params)}) capture_client_event("client.delete_all", self, {"keys": list(kwargs.keys())})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -257,7 +263,7 @@ class MemoryClient:
""" """
response = self.client.get(f"/v1/memories/{memory_id}/history/") response = self.client.get(f"/v1/memories/{memory_id}/history/")
response.raise_for_status() response.raise_for_status()
capture_client_event("client.history", self) capture_client_event("client.history", self, {"memory_id": memory_id})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -390,14 +396,16 @@ class AsyncMemoryClient:
payload = self.sync_client._prepare_payload(messages, kwargs) payload = self.sync_client._prepare_payload(messages, kwargs)
response = await self.async_client.post("/v1/memories/", json=payload) response = await self.async_client.post("/v1/memories/", json=payload)
response.raise_for_status() response.raise_for_status()
capture_client_event("async_client.add", self.sync_client) if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event("async_client.add", self.sync_client, {"keys": list(kwargs.keys())})
return response.json() return response.json()
@api_error_handler @api_error_handler
async def get(self, memory_id: str) -> Dict[str, Any]: async def get(self, memory_id: str) -> Dict[str, Any]:
response = await self.async_client.get(f"/v1/memories/{memory_id}/") response = await self.async_client.get(f"/v1/memories/{memory_id}/")
response.raise_for_status() response.raise_for_status()
capture_client_event("async_client.get", self.sync_client) capture_client_event("async_client.get", self.sync_client, {"memory_id": memory_id})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -408,8 +416,10 @@ class AsyncMemoryClient:
elif version == "v2": elif version == "v2":
response = await self.async_client.post(f"/{version}/memories/", json=params) response = await self.async_client.post(f"/{version}/memories/", json=params)
response.raise_for_status() response.raise_for_status()
if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event( capture_client_event(
"async_client.get_all", self.sync_client, {"filters": len(params), "limit": kwargs.get("limit", 100)} "async_client.get_all", self.sync_client, {"api_version": version, "keys": list(kwargs.keys())}
) )
return response.json() return response.json()
@@ -419,21 +429,25 @@ class AsyncMemoryClient:
payload.update(self.sync_client._prepare_params(kwargs)) payload.update(self.sync_client._prepare_params(kwargs))
response = await self.async_client.post(f"/{version}/memories/search/", json=payload) response = await self.async_client.post(f"/{version}/memories/search/", json=payload)
response.raise_for_status() response.raise_for_status()
capture_client_event("async_client.search", self.sync_client, {"limit": kwargs.get("limit", 100)}) if "metadata" in kwargs:
del kwargs["metadata"]
capture_client_event(
"async_client.search", self.sync_client, {"api_version": version, "keys": list(kwargs.keys())}
)
return response.json() return response.json()
@api_error_handler @api_error_handler
async def update(self, memory_id: str, data: str) -> Dict[str, Any]: async def update(self, memory_id: str, data: str) -> Dict[str, Any]:
response = await self.async_client.put(f"/v1/memories/{memory_id}/", json={"text": data}) response = await self.async_client.put(f"/v1/memories/{memory_id}/", json={"text": data})
response.raise_for_status() response.raise_for_status()
capture_client_event("async_client.update", self.sync_client) capture_client_event("async_client.update", self.sync_client, {"memory_id": memory_id})
return response.json() return response.json()
@api_error_handler @api_error_handler
async def delete(self, memory_id: str) -> Dict[str, Any]: async def delete(self, memory_id: str) -> Dict[str, Any]:
response = await self.async_client.delete(f"/v1/memories/{memory_id}/") response = await self.async_client.delete(f"/v1/memories/{memory_id}/")
response.raise_for_status() response.raise_for_status()
capture_client_event("async_client.delete", self.sync_client) capture_client_event("async_client.delete", self.sync_client, {"memory_id": memory_id})
return response.json() return response.json()
@api_error_handler @api_error_handler
@@ -441,14 +455,14 @@ class AsyncMemoryClient:
params = self.sync_client._prepare_params(kwargs) params = self.sync_client._prepare_params(kwargs)
response = await self.async_client.delete("/v1/memories/", params=params) response = await self.async_client.delete("/v1/memories/", params=params)
response.raise_for_status() response.raise_for_status()
capture_client_event("async_client.delete_all", self.sync_client, {"params": len(params)}) capture_client_event("async_client.delete_all", self.sync_client, {"keys": list(kwargs.keys())})
return response.json() return response.json()
@api_error_handler @api_error_handler
async def history(self, memory_id: str) -> List[Dict[str, Any]]: async def history(self, memory_id: str) -> List[Dict[str, Any]]:
response = await self.async_client.get(f"/v1/memories/{memory_id}/history/") response = await self.async_client.get(f"/v1/memories/{memory_id}/history/")
response.raise_for_status() response.raise_for_status()
capture_client_event("async_client.history", self.sync_client) capture_client_event("async_client.history", self.sync_client, {"memory_id": memory_id})
return response.json() return response.json()
@api_error_handler @api_error_handler

2
mem0/configs/vector_stores/chroma.py
View File

@@ -1,5 +1,3 @@
import subprocess
import sys
from typing import Any, ClassVar, Dict, Optional from typing import Any, ClassVar, Dict, Optional
from pydantic import BaseModel, Field, model_validator from pydantic import BaseModel, Field, model_validator

29
mem0/memory/main.py
View File

@@ -37,11 +37,11 @@ class Memory(MemoryBase):
self.llm = LlmFactory.create(self.config.llm.provider, self.config.llm.config) self.llm = LlmFactory.create(self.config.llm.provider, self.config.llm.config)
self.db = SQLiteManager(self.config.history_db_path) self.db = SQLiteManager(self.config.history_db_path)
self.collection_name = self.config.vector_store.config.collection_name self.collection_name = self.config.vector_store.config.collection_name
self.version = self.config.version self.api_version = self.config.version
self.enable_graph = False self.enable_graph = False
if self.version == "v1.1" and self.config.graph_store.config: if self.api_version == "v1.1" and self.config.graph_store.config:
from mem0.memory.graph_memory import MemoryGraph from mem0.memory.graph_memory import MemoryGraph
self.graph = MemoryGraph(self.config) self.graph = MemoryGraph(self.config)
@@ -119,7 +119,7 @@ class Memory(MemoryBase):
vector_store_result = future1.result() vector_store_result = future1.result()
graph_result = future2.result() graph_result = future2.result()
if self.version == "v1.1": if self.api_version == "v1.1":
return { return {
"results": vector_store_result, "results": vector_store_result,
"relations": graph_result, "relations": graph_result,
@@ -226,13 +226,13 @@ class Memory(MemoryBase):
except Exception as e: except Exception as e:
logging.error(f"Error in new_memories_with_actions: {e}") logging.error(f"Error in new_memories_with_actions: {e}")
capture_event("mem0.add", self) capture_event("mem0.add", self, {"version": self.api_version, "keys": list(filters.keys())})
return returned_memories return returned_memories
def _add_to_graph(self, messages, filters): def _add_to_graph(self, messages, filters):
added_entities = [] added_entities = []
if self.version == "v1.1" and self.enable_graph: if self.api_version == "v1.1" and self.enable_graph:
if filters["user_id"]: if filters["user_id"]:
self.graph.user_id = filters["user_id"] self.graph.user_id = filters["user_id"]
elif filters["agent_id"]: elif filters["agent_id"]:
@@ -305,13 +305,13 @@ class Memory(MemoryBase):
if run_id: if run_id:
filters["run_id"] = run_id filters["run_id"] = run_id
capture_event("mem0.get_all", self, {"filters": len(filters), "limit": limit}) capture_event("mem0.get_all", self, {"limit": limit, "keys": list(filters.keys())})
with concurrent.futures.ThreadPoolExecutor() as executor: with concurrent.futures.ThreadPoolExecutor() as executor:
future_memories = executor.submit(self._get_all_from_vector_store, filters, limit) future_memories = executor.submit(self._get_all_from_vector_store, filters, limit)
future_graph_entities = ( future_graph_entities = (
executor.submit(self.graph.get_all, filters, limit) executor.submit(self.graph.get_all, filters, limit)
if self.version == "v1.1" and self.enable_graph if self.api_version == "v1.1" and self.enable_graph
else None else None
) )
@@ -322,7 +322,7 @@ class Memory(MemoryBase):
all_memories = future_memories.result() all_memories = future_memories.result()
graph_entities = future_graph_entities.result() if future_graph_entities else None graph_entities = future_graph_entities.result() if future_graph_entities else None
if self.version == "v1.1": if self.api_version == "v1.1":
if self.enable_graph: if self.enable_graph:
return {"results": all_memories, "relations": graph_entities} return {"results": all_memories, "relations": graph_entities}
else: else:
@@ -398,14 +398,14 @@ class Memory(MemoryBase):
capture_event( capture_event(
"mem0.search", "mem0.search",
self, self,
{"filters": len(filters), "limit": limit, "version": self.version}, {"limit": limit, "version": self.api_version, "keys": list(filters.keys())},
) )
with concurrent.futures.ThreadPoolExecutor() as executor: with concurrent.futures.ThreadPoolExecutor() as executor:
future_memories = executor.submit(self._search_vector_store, query, filters, limit) future_memories = executor.submit(self._search_vector_store, query, filters, limit)
future_graph_entities = ( future_graph_entities = (
executor.submit(self.graph.search, query, filters, limit) executor.submit(self.graph.search, query, filters, limit)
if self.version == "v1.1" and self.enable_graph if self.api_version == "v1.1" and self.enable_graph
else None else None
) )
@@ -416,7 +416,7 @@ class Memory(MemoryBase):
original_memories = future_memories.result() original_memories = future_memories.result()
graph_entities = future_graph_entities.result() if future_graph_entities else None graph_entities = future_graph_entities.result() if future_graph_entities else None
if self.version == "v1.1": if self.api_version == "v1.1":
if self.enable_graph: if self.enable_graph:
return {"results": original_memories, "relations": graph_entities} return {"results": original_memories, "relations": graph_entities}
else: else:
@@ -518,14 +518,14 @@ class Memory(MemoryBase):
"At least one filter is required to delete all memories. If you want to delete all memories, use the `reset()` method." "At least one filter is required to delete all memories. If you want to delete all memories, use the `reset()` method."
) )
capture_event("mem0.delete_all", self, {"filters": len(filters)}) capture_event("mem0.delete_all", self, {"keys": list(filters.keys())})
memories = self.vector_store.list(filters=filters)[0] memories = self.vector_store.list(filters=filters)[0]
for memory in memories: for memory in memories:
self._delete_memory(memory.id) self._delete_memory(memory.id)
logger.info(f"Deleted {len(memories)} memories") logger.info(f"Deleted {len(memories)} memories")
if self.version == "v1.1" and self.enable_graph: if self.api_version == "v1.1" and self.enable_graph:
self.graph.delete_all(filters) self.graph.delete_all(filters)
return {"message": "Memories deleted successfully!"} return {"message": "Memories deleted successfully!"}
@@ -561,6 +561,7 @@ class Memory(MemoryBase):
payloads=[metadata], payloads=[metadata],
) )
self.db.add_history(memory_id, None, data, "ADD", created_at=metadata["created_at"]) self.db.add_history(memory_id, None, data, "ADD", created_at=metadata["created_at"])
capture_event("mem0._create_memory", self, {"memory_id": memory_id})
return memory_id return memory_id
def _update_memory(self, memory_id, data, existing_embeddings, metadata=None): def _update_memory(self, memory_id, data, existing_embeddings, metadata=None):
@@ -603,6 +604,7 @@ class Memory(MemoryBase):
created_at=new_metadata["created_at"], created_at=new_metadata["created_at"],
updated_at=new_metadata["updated_at"], updated_at=new_metadata["updated_at"],
) )
capture_event("mem0._update_memory", self, {"memory_id": memory_id})
return memory_id return memory_id
def _delete_memory(self, memory_id): def _delete_memory(self, memory_id):
@@ -611,6 +613,7 @@ class Memory(MemoryBase):
prev_value = existing_memory.payload["data"] prev_value = existing_memory.payload["data"]
self.vector_store.delete(vector_id=memory_id) self.vector_store.delete(vector_id=memory_id)
self.db.add_history(memory_id, prev_value, None, "DELETE", is_deleted=1) self.db.add_history(memory_id, prev_value, None, "DELETE", is_deleted=1)
capture_event("mem0._delete_memory", self, {"memory_id": memory_id})
return memory_id return memory_id
def reset(self): def reset(self):

2
mem0/memory/telemetry.py
View File

@@ -67,7 +67,7 @@ def capture_event(event_name, memory_instance, additional_data=None):
"vector_store": f"{memory_instance.vector_store.__class__.__module__}.{memory_instance.vector_store.__class__.__name__}", "vector_store": f"{memory_instance.vector_store.__class__.__module__}.{memory_instance.vector_store.__class__.__name__}",
"llm": f"{memory_instance.llm.__class__.__module__}.{memory_instance.llm.__class__.__name__}", "llm": f"{memory_instance.llm.__class__.__module__}.{memory_instance.llm.__class__.__name__}",
"embedding_model": f"{memory_instance.embedding_model.__class__.__module__}.{memory_instance.embedding_model.__class__.__name__}", "embedding_model": f"{memory_instance.embedding_model.__class__.__module__}.{memory_instance.embedding_model.__class__.__name__}",
"function": f"{memory_instance.__class__.__module__}.{memory_instance.__class__.__name__}.{memory_instance.version}", "function": f"{memory_instance.__class__.__module__}.{memory_instance.__class__.__name__}.{memory_instance.api_version}",
} }
if additional_data: if additional_data:
event_data.update(additional_data) event_data.update(additional_data)

10
tests/embeddings/test_azure_openai_embeddings.py
View File

@@ -33,11 +33,7 @@ def test_embed_text(mock_openai_client):
@pytest.mark.parametrize( @pytest.mark.parametrize(
"default_headers, expected_header", "default_headers, expected_header",
[ [(None, None), ({"Test": "test_value"}, "test_value"), ({}, None)],
(None, None),
({"Test": "test_value"}, "test_value"),
({}, None)
],
) )
def test_embed_text_with_default_headers(default_headers, expected_header): def test_embed_text_with_default_headers(default_headers, expected_header):
config = BaseEmbedderConfig( config = BaseEmbedderConfig(
@@ -47,8 +43,8 @@ def test_embed_text_with_default_headers(default_headers, expected_header):
"api_version": "test_version", "api_version": "test_version",
"azure_endpoint": "test_endpoint", "azure_endpoint": "test_endpoint",
"azuer_deployment": "test_deployment", "azuer_deployment": "test_deployment",
"default_headers": default_headers "default_headers": default_headers,
} },
) )
embedder = AzureOpenAIEmbedding(config) embedder = AzureOpenAIEmbedding(config)
assert embedder.client.api_key == "test" assert embedder.client.api_key == "test"

15
tests/embeddings/test_gemini.py
View File

@@ -12,17 +12,11 @@ def mock_genai():
@pytest.fixture @pytest.fixture
def config(): def config():
return BaseEmbedderConfig( return BaseEmbedderConfig(api_key="dummy_api_key", model="test_model")
api_key="dummy_api_key",
model="test_model"
)
def test_embed_query(mock_genai, config): def test_embed_query(mock_genai, config):
mock_embedding_response = {"embedding": [0.1, 0.2, 0.3, 0.4]}
mock_embedding_response = {
'embedding': [0.1, 0.2, 0.3, 0.4]
}
mock_genai.return_value = mock_embedding_response mock_genai.return_value = mock_embedding_response
embedder = GoogleGenAIEmbedding(config) embedder = GoogleGenAIEmbedding(config)
@@ -31,7 +25,4 @@ def test_embed_query(mock_genai, config):
embedding = embedder.embed(text) embedding = embedder.embed(text)
assert embedding == [0.1, 0.2, 0.3, 0.4] assert embedding == [0.1, 0.2, 0.3, 0.4]
mock_genai.assert_called_once_with( mock_genai.assert_called_once_with(model="test_model", content="Hello, world!")
model="test_model",
content="Hello, world!"
)

31
tests/embeddings/test_vertexai_embeddings.py
View File

@@ -1,7 +1,6 @@
import pytest import pytest
from unittest.mock import Mock, patch from unittest.mock import Mock, patch
from mem0.embeddings.vertexai import VertexAIEmbedding from mem0.embeddings.vertexai import VertexAIEmbedding
from mem0.configs.embeddings.base import BaseEmbedderConfig
@pytest.fixture @pytest.fixture
@@ -35,15 +34,11 @@ def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_co
embedder = VertexAIEmbedding(config) embedder = VertexAIEmbedding(config)
mock_embedding = Mock(values=[0.1, 0.2, 0.3]) mock_embedding = Mock(values=[0.1, 0.2, 0.3])
mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [ mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
mock_embedding
]
result = embedder.embed("Hello world") embedder.embed("Hello world")
mock_text_embedding_model.from_pretrained.assert_called_once_with( mock_text_embedding_model.from_pretrained.assert_called_once_with("text-embedding-004")
"text-embedding-004"
)
mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with( mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with(
texts=["Hello world"], output_dimensionality=256 texts=["Hello world"], output_dimensionality=256
) )
@@ -60,15 +55,11 @@ def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_con
embedder = VertexAIEmbedding(config) embedder = VertexAIEmbedding(config)
mock_embedding = Mock(values=[0.4, 0.5, 0.6]) mock_embedding = Mock(values=[0.4, 0.5, 0.6])
mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [ mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
mock_embedding
]
result = embedder.embed("Test embedding") result = embedder.embed("Test embedding")
mock_text_embedding_model.from_pretrained.assert_called_with( mock_text_embedding_model.from_pretrained.assert_called_with("custom-embedding-model")
"custom-embedding-model"
)
mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with( mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with(
texts=["Test embedding"], output_dimensionality=512 texts=["Test embedding"], output_dimensionality=512
) )
@@ -93,16 +84,12 @@ def test_missing_credentials(mock_os, mock_text_embedding_model, mock_config):
config = mock_config() config = mock_config()
with pytest.raises( with pytest.raises(ValueError, match="Google application credentials JSON is not provided"):
ValueError, match="Google application credentials JSON is not provided"
):
VertexAIEmbedding(config) VertexAIEmbedding(config)
@patch("mem0.embeddings.vertexai.TextEmbeddingModel") @patch("mem0.embeddings.vertexai.TextEmbeddingModel")
def test_embed_with_different_dimensions( def test_embed_with_different_dimensions(mock_text_embedding_model, mock_os_environ, mock_config):
mock_text_embedding_model, mock_os_environ, mock_config
):
mock_config.vertex_credentials_json = "/path/to/credentials.json" mock_config.vertex_credentials_json = "/path/to/credentials.json"
mock_config.return_value.embedding_dims = 1024 mock_config.return_value.embedding_dims = 1024
@@ -110,9 +97,7 @@ def test_embed_with_different_dimensions(
embedder = VertexAIEmbedding(config) embedder = VertexAIEmbedding(config)
mock_embedding = Mock(values=[0.1] * 1024) mock_embedding = Mock(values=[0.1] * 1024)
mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [ mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
mock_embedding
]
result = embedder.embed("Large embedding test") result = embedder.embed("Large embedding test")

61
tests/llms/test_gemini_llm.py
View File

@@ -33,19 +33,19 @@ def test_generate_response_without_tools(mock_gemini_client: Mock):
response = llm.generate_response(messages) response = llm.generate_response(messages)
mock_gemini_client.generate_content.assert_called_once_with( mock_gemini_client.generate_content.assert_called_once_with(
contents = [ contents=[
{"parts": "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: You are a helpful assistant.", "role": "user"}, {"parts": "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: You are a helpful assistant.", "role": "user"},
{"parts": "Hello, how are you?", "role": "user"} {"parts": "Hello, how are you?", "role": "user"},
], ],
generation_config = GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0), generation_config=GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0),
tools = None, tools=None,
tool_config = content_types.to_tool_config( tool_config=content_types.to_tool_config(
{"function_calling_config": {"function_calling_config": {"mode": "auto", "allowed_function_names": None}}
{"mode": 'auto', "allowed_function_names": None} ),
})
) )
assert response == "I'm doing well, thank you for asking!" assert response == "I'm doing well, thank you for asking!"
def test_generate_response_with_tools(mock_gemini_client: Mock): def test_generate_response_with_tools(mock_gemini_client: Mock):
config = BaseLlmConfig(model="gemini-1.5-flash-latest", temperature=0.7, max_tokens=100, top_p=1.0) config = BaseLlmConfig(model="gemini-1.5-flash-latest", temperature=0.7, max_tokens=100, top_p=1.0)
llm = GeminiLLM(config) llm = GeminiLLM(config)
@@ -74,13 +74,13 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
mock_part = Mock() mock_part = Mock()
mock_part.function_call = mock_tool_call mock_part.function_call = mock_tool_call
mock_part.text="I've added the memory for you." mock_part.text = "I've added the memory for you."
mock_content = Mock() mock_content = Mock()
mock_content.parts=[mock_part] mock_content.parts = [mock_part]
mock_message = Mock() mock_message = Mock()
mock_message.content=mock_content mock_message.content = mock_content
mock_response = Mock(candidates=[mock_message]) mock_response = Mock(candidates=[mock_message])
mock_gemini_client.generate_content.return_value = mock_response mock_gemini_client.generate_content.return_value = mock_response
@@ -88,28 +88,29 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
response = llm.generate_response(messages, tools=tools) response = llm.generate_response(messages, tools=tools)
mock_gemini_client.generate_content.assert_called_once_with( mock_gemini_client.generate_content.assert_called_once_with(
contents = [ contents=[
{"parts": "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: You are a helpful assistant.", "role": "user"}, {"parts": "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: You are a helpful assistant.", "role": "user"},
{"parts": "Add a new memory: Today is a sunny day.", "role": "user"} {"parts": "Add a new memory: Today is a sunny day.", "role": "user"},
], ],
generation_config = GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0), generation_config=GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0),
tools = [ tools=[
{ {
"function_declarations": [{ "function_declarations": [
"name": "add_memory", {
"description": "Add a memory", "name": "add_memory",
"parameters": { "description": "Add a memory",
"type": "object", "parameters": {
"properties": {"data": {"type": "string", "description": "Data to add to memory"}}, "type": "object",
"required": ["data"] "properties": {"data": {"type": "string", "description": "Data to add to memory"}},
"required": ["data"],
},
} }
}] ]
} }
], ],
tool_config = content_types.to_tool_config( tool_config=content_types.to_tool_config(
{"function_calling_config": {"function_calling_config": {"mode": "auto", "allowed_function_names": None}}
{"mode": 'auto', "allowed_function_names": None} ),
})
) )
assert response["content"] == "I've added the memory for you." assert response["content"] == "I've added the memory for you."

5
tests/llms/test_openai.py
View File

@@ -31,8 +31,9 @@ def test_openai_llm_base_url():
# case3: with config.openai_base_url # case3: with config.openai_base_url
config_base_url = "https://api.config.com/v1" config_base_url = "https://api.config.com/v1"
config = BaseLlmConfig(model="gpt-4o", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key", config = BaseLlmConfig(
openai_base_url=config_base_url) model="gpt-4o", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key", openai_base_url=config_base_url
)
llm = OpenAILLM(config) llm = OpenAILLM(config)
# Note: openai client will parse the raw base_url into a URL object, which will have a trailing slash # Note: openai client will parse the raw base_url into a URL object, which will have a trailing slash
assert str(llm.client.base_url) == config_base_url + "/" assert str(llm.client.base_url) == config_base_url + "/"

6
tests/test_main.py
View File

@@ -135,7 +135,9 @@ def test_update(memory_instance):
result = memory_instance.update("test_id", "Updated memory") result = memory_instance.update("test_id", "Updated memory")
memory_instance._update_memory.assert_called_once_with("test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}) memory_instance._update_memory.assert_called_once_with(
"test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}
)
assert result["message"] == "Memory updated successfully!" assert result["message"] == "Memory updated successfully!"
@@ -177,7 +179,6 @@ def test_reset(memory_instance):
memory_instance.db.reset = Mock() memory_instance.db.reset = Mock()
with patch.object(VectorStoreFactory, "create", return_value=Mock()) as mock_create: with patch.object(VectorStoreFactory, "create", return_value=Mock()) as mock_create:
memory_instance.reset() memory_instance.reset()
initial_vector_store.delete_col.assert_called_once() initial_vector_store.delete_col.assert_called_once()
@@ -186,6 +187,7 @@ def test_reset(memory_instance):
memory_instance.config.vector_store.provider, memory_instance.config.vector_store.config memory_instance.config.vector_store.provider, memory_instance.config.vector_store.config
) )
@pytest.mark.parametrize( @pytest.mark.parametrize(
"version, enable_graph, expected_result", "version, enable_graph, expected_result",
[ [

22
tests/vector_stores/test_chroma.py
View File

@@ -1,6 +1,6 @@
from unittest.mock import Mock, patch from unittest.mock import Mock, patch
import pytest import pytest
from mem0.vector_stores.chroma import ChromaDB, OutputData from mem0.vector_stores.chroma import ChromaDB
@pytest.fixture @pytest.fixture
@@ -12,13 +12,9 @@ def mock_chromadb_client():
@pytest.fixture @pytest.fixture
def chromadb_instance(mock_chromadb_client): def chromadb_instance(mock_chromadb_client):
mock_collection = Mock() mock_collection = Mock()
mock_chromadb_client.return_value.get_or_create_collection.return_value = ( mock_chromadb_client.return_value.get_or_create_collection.return_value = mock_collection
mock_collection
)
return ChromaDB( return ChromaDB(collection_name="test_collection", client=mock_chromadb_client.return_value)
collection_name="test_collection", client=mock_chromadb_client.return_value
)
def test_insert_vectors(chromadb_instance, mock_chromadb_client): def test_insert_vectors(chromadb_instance, mock_chromadb_client):
@@ -28,9 +24,7 @@ def test_insert_vectors(chromadb_instance, mock_chromadb_client):
chromadb_instance.insert(vectors=vectors, payloads=payloads, ids=ids) chromadb_instance.insert(vectors=vectors, payloads=payloads, ids=ids)
chromadb_instance.collection.add.assert_called_once_with( chromadb_instance.collection.add.assert_called_once_with(ids=ids, embeddings=vectors, metadatas=payloads)
ids=ids, embeddings=vectors, metadatas=payloads
)
def test_search_vectors(chromadb_instance, mock_chromadb_client): def test_search_vectors(chromadb_instance, mock_chromadb_client):
@@ -44,9 +38,7 @@ def test_search_vectors(chromadb_instance, mock_chromadb_client):
query = [[0.1, 0.2, 0.3]] query = [[0.1, 0.2, 0.3]]
results = chromadb_instance.search(query=query, limit=2) results = chromadb_instance.search(query=query, limit=2)
chromadb_instance.collection.query.assert_called_once_with( chromadb_instance.collection.query.assert_called_once_with(query_embeddings=query, where=None, n_results=2)
query_embeddings=query, where=None, n_results=2
)
print(results, type(results)) print(results, type(results))
assert len(results) == 2 assert len(results) == 2
@@ -68,9 +60,7 @@ def test_update_vector(chromadb_instance):
new_vector = [0.7, 0.8, 0.9] new_vector = [0.7, 0.8, 0.9]
new_payload = {"name": "updated_vector"} new_payload = {"name": "updated_vector"}
chromadb_instance.update( chromadb_instance.update(vector_id=vector_id, vector=new_vector, payload=new_payload)
vector_id=vector_id, vector=new_vector, payload=new_payload
)
chromadb_instance.collection.update.assert_called_once_with( chromadb_instance.collection.update.assert_called_once_with(
ids=vector_id, embeddings=new_vector, metadatas=new_payload ids=vector_id, embeddings=new_vector, metadatas=new_payload

26
tests/vector_stores/test_qdrant.py
View File

@@ -1,5 +1,5 @@
import unittest import unittest
from unittest.mock import MagicMock, patch from unittest.mock import MagicMock
import uuid import uuid
from qdrant_client import QdrantClient from qdrant_client import QdrantClient
from qdrant_client.models import ( from qdrant_client.models import (
@@ -51,9 +51,7 @@ class TestQdrant(unittest.TestCase):
def test_search(self): def test_search(self):
query_vector = [0.1, 0.2] query_vector = [0.1, 0.2]
self.client_mock.search.return_value = [ self.client_mock.search.return_value = [{"id": str(uuid.uuid4()), "score": 0.95, "payload": {"key": "value"}}]
{"id": str(uuid.uuid4()), "score": 0.95, "payload": {"key": "value"}}
]
results = self.qdrant.search(query=query_vector, limit=1) results = self.qdrant.search(query=query_vector, limit=1)
@@ -83,9 +81,7 @@ class TestQdrant(unittest.TestCase):
updated_vector = [0.2, 0.3] updated_vector = [0.2, 0.3]
updated_payload = {"key": "updated_value"} updated_payload = {"key": "updated_value"}
self.qdrant.update( self.qdrant.update(vector_id=vector_id, vector=updated_vector, payload=updated_payload)
vector_id=vector_id, vector=updated_vector, payload=updated_payload
)
self.client_mock.upsert.assert_called_once() self.client_mock.upsert.assert_called_once()
point = self.client_mock.upsert.call_args[1]["points"][0] point = self.client_mock.upsert.call_args[1]["points"][0]
@@ -95,9 +91,7 @@ class TestQdrant(unittest.TestCase):
def test_get(self): def test_get(self):
vector_id = str(uuid.uuid4()) vector_id = str(uuid.uuid4())
self.client_mock.retrieve.return_value = [ self.client_mock.retrieve.return_value = [{"id": vector_id, "payload": {"key": "value"}}]
{"id": vector_id, "payload": {"key": "value"}}
]
result = self.qdrant.get(vector_id=vector_id) result = self.qdrant.get(vector_id=vector_id)
@@ -108,23 +102,17 @@ class TestQdrant(unittest.TestCase):
self.assertEqual(result["payload"], {"key": "value"}) self.assertEqual(result["payload"], {"key": "value"})
def test_list_cols(self): def test_list_cols(self):
self.client_mock.get_collections.return_value = MagicMock( self.client_mock.get_collections.return_value = MagicMock(collections=[{"name": "test_collection"}])
collections=[{"name": "test_collection"}]
)
result = self.qdrant.list_cols() result = self.qdrant.list_cols()
self.assertEqual(result.collections[0]["name"], "test_collection") self.assertEqual(result.collections[0]["name"], "test_collection")
def test_delete_col(self): def test_delete_col(self):
self.qdrant.delete_col() self.qdrant.delete_col()
self.client_mock.delete_collection.assert_called_once_with( self.client_mock.delete_collection.assert_called_once_with(collection_name="test_collection")
collection_name="test_collection"
)
def test_col_info(self): def test_col_info(self):
self.qdrant.col_info() self.qdrant.col_info()
self.client_mock.get_collection.assert_called_once_with( self.client_mock.get_collection.assert_called_once_with(collection_name="test_collection")
collection_name="test_collection"
)
def tearDown(self): def tearDown(self):
del self.qdrant del self.qdrant