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

Fix CI issues related to missing dependency (#3096)

Browse Source
This commit is contained in:
Deshraj Yadav
2025-07-03 18:52:50 -07:00
committed by GitHub
parent 2c496e6376
commit 7484eed4b2
32 changed files with 6150 additions and 828 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
tests/embeddings/test_gemini_emeddings.py
View File

@@ -1,4 +1,4 @@
from unittest.mock import patch
from unittest.mock import patch, ANY
import pytest
@@ -8,8 +8,10 @@ from mem0.embeddings.gemini import GoogleGenAIEmbedding
@pytest.fixture
def mock_genai():
with patch("mem0.embeddings.gemini.genai.embed_content") as mock_genai:
yield mock_genai
with patch("mem0.embeddings.gemini.genai.Client") as mock_client_class:
mock_client = mock_client_class.return_value
mock_client.models.embed_content.return_value = None
yield mock_client.models.embed_content
@pytest.fixture
@@ -18,7 +20,9 @@ def config():
def test_embed_query(mock_genai, config):
mock_embedding_response = {"embedding": [0.1, 0.2, 0.3, 0.4]}
mock_embedding_response = type('Response', (), {
'embeddings': [type('Embedding', (), {'values': [0.1, 0.2, 0.3, 0.4]})]
})()
mock_genai.return_value = mock_embedding_response
embedder = GoogleGenAIEmbedding(config)
@@ -27,10 +31,11 @@ def test_embed_query(mock_genai, config):
embedding = embedder.embed(text)
assert embedding == [0.1, 0.2, 0.3, 0.4]
mock_genai.assert_called_once_with(model="test_model", content="Hello, world!", output_dimensionality=786)
mock_genai.assert_called_once_with(model="test_model", contents="Hello, world!", config=ANY)
def test_embed_returns_empty_list_if_none(mock_genai, config):
mock_genai.return_value = None
mock_genai.return_value = type('Response', (), {'embeddings': [type('Embedding', (), {'values': []})]})()
embedder = GoogleGenAIEmbedding(config)
result = embedder.embed("test")
@@ -47,10 +52,10 @@ def test_embed_raises_on_error(mock_genai, config):
with pytest.raises(RuntimeError, match="Embedding failed"):
embedder.embed("some input")
def test_config_initialization(config):
embedder = GoogleGenAIEmbedding(config)
assert embedder.config.api_key == "dummy_api_key"
assert embedder.config.model == "test_model"
assert embedder.config.embedding_dims == 786

131
tests/llms/test_gemini.py
View File

@@ -9,7 +9,7 @@ from mem0.llms.gemini import GeminiLLM
@pytest.fixture
def mock_gemini_client():
with patch("mem0.llms.gemini.genai") as mock_client_class:
with patch("mem0.llms.gemini.genai.Client") as mock_client_class:
mock_client = Mock()
mock_client_class.return_value = mock_client
yield mock_client
@@ -24,43 +24,30 @@ def test_generate_response_without_tools(mock_gemini_client: Mock):
]
mock_part = Mock(text="I'm doing well, thank you for asking!")
mock_embedding = Mock()
mock_embedding.values = [0.1, 0.2, 0.3]
mock_response = Mock()
mock_response.candidates = [Mock()]
mock_response.candidates[0].content.parts = [Mock()]
mock_response.candidates[0].content.parts[0].text = "I'm doing well, thank you for asking!"
mock_content = Mock(parts=[mock_part])
mock_candidate = Mock(content=mock_content)
mock_response = Mock(candidates=[mock_candidate])
mock_gemini_client.models.generate_content.return_value = mock_response
mock_content = Mock(parts=[mock_part])
mock_message = Mock(content=mock_content)
mock_response = Mock(candidates=[mock_message])
mock_gemini_client.generate_content.return_value = mock_response
response = llm.generate_response(messages)
mock_gemini_client.generate_content.assert_called_once_with(
contents=[
{"parts": "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: You are a helpful assistant.", "role": "user"},
{"parts": "Hello, how are you?", "role": "user"},
],
config=types.GenerateContentConfig(
temperature=0.7,
max_output_tokens=100,
top_p=1.0,
tools=None,
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
allowed_function_names=None,
mode="auto"
)
)
) )
assert response == "I'm doing well, thank you for asking!"
# Check the actual call - system instruction is now in config
mock_gemini_client.models.generate_content.assert_called_once()
call_args = mock_gemini_client.models.generate_content.call_args
# Verify model and contents
assert call_args.kwargs['model'] == "gemini-2.0-flash-latest"
assert len(call_args.kwargs['contents']) == 1 # Only user message
# Verify config has system instruction
config_arg = call_args.kwargs['config']
assert config_arg.system_instruction == "You are a helpful assistant."
assert config_arg.temperature == 0.7
assert config_arg.max_output_tokens == 100
assert config_arg.top_p == 1.0
assert response == "I'm doing well, thank you for asking!"
def test_generate_response_with_tools(mock_gemini_client: Mock):
@@ -89,64 +76,42 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
mock_tool_call.name = "add_memory"
mock_tool_call.args = {"data": "Today is a sunny day."}
mock_part = Mock()
mock_part.function_call = mock_tool_call
mock_part.text = "I've added the memory for you."
# Create mock parts with both text and function_call
mock_text_part = Mock()
mock_text_part.text = "I've added the memory for you."
mock_text_part.function_call = None
mock_func_part = Mock()
mock_func_part.text = None
mock_func_part.function_call = mock_tool_call
mock_content = Mock()
mock_content.parts = [mock_part]
mock_content.parts = [mock_text_part, mock_func_part]
mock_message = Mock()
mock_message.content = mock_content
mock_candidate = Mock()
mock_candidate.content = mock_content
mock_response = Mock(candidates=[mock_message])
mock_gemini_client.generate_content.return_value = mock_response
mock_response = Mock(candidates=[mock_candidate])
mock_gemini_client.models.generate_content.return_value = mock_response
response = llm.generate_response(messages, tools=tools)
mock_gemini_client.generate_content.assert_called_once_with(
contents=[
{
"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"
},
],
config=types.GenerateContentConfig(
temperature=0.7,
max_output_tokens=100,
top_p=1.0,
tools=[
types.Tool(
function_declarations=[
types.FunctionDeclaration(
name="add_memory",
description="Add a memory",
parameters={
"type": "object",
"properties": {
"data": {
"type": "string",
"description": "Data to add to memory"
}
},
"required": ["data"]
}
)
]
)
],
tool_config=types.ToolConfig(
function_calling_config=types.FunctionCallingConfig(
allowed_function_names=None,
mode="auto"
)
)
)
)
# Check the actual call
mock_gemini_client.models.generate_content.assert_called_once()
call_args = mock_gemini_client.models.generate_content.call_args
# Verify model and contents
assert call_args.kwargs['model'] == "gemini-1.5-flash-latest"
assert len(call_args.kwargs['contents']) == 1 # Only user message
# Verify config has system instruction and tools
config_arg = call_args.kwargs['config']
assert config_arg.system_instruction == "You are a helpful assistant."
assert config_arg.temperature == 0.7
assert config_arg.max_output_tokens == 100
assert config_arg.top_p == 1.0
assert len(config_arg.tools) == 1
assert config_arg.tool_config.function_calling_config.mode == types.FunctionCallingConfigMode.AUTO
assert response["content"] == "I've added the memory for you."
assert len(response["tool_calls"]) == 1

3
tests/llms/test_lm_studio.py
View File

@@ -43,6 +43,7 @@ def test_generate_response_without_tools(mock_lm_studio_client):
assert response == "I'm doing well, thank you for asking!"
def test_generate_response_specifying_response_format(mock_lm_studio_client):
config = BaseLlmConfig(
model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
@@ -68,4 +69,4 @@ def test_generate_response_specifying_response_format(mock_lm_studio_client):
response_format={"type": "json_schema"},
)
assert response == "I'm doing well, thank you for asking!"
assert response == "I'm doing well, thank you for asking!"

8
tests/llms/test_vllm.py
View File

@@ -71,7 +71,13 @@ def test_generate_response_with_tools(mock_vllm_client):
response = llm.generate_response(messages, tools=tools)
mock_vllm_client.chat.completions.create.assert_called_once_with(
model="Qwen/Qwen2.5-32B-Instruct", messages=messages, temperature=0.7, max_tokens=100, top_p=1.0, tools=tools, tool_choice="auto"
model="Qwen/Qwen2.5-32B-Instruct",
messages=messages,
temperature=0.7,
max_tokens=100,
top_p=1.0,
tools=tools,
tool_choice="auto",
)
assert response["content"] == "I've added the memory for you."

2
tests/test_main.py
View File

@@ -253,10 +253,10 @@ def test_get_all(memory_instance, version, enable_graph, expected_result):
def test_custom_prompts(memory_custom_instance):
messages = [{"role": "user", "content": "Test message"}]
from mem0.embeddings.mock import MockEmbeddings
memory_custom_instance.llm.generate_response = Mock()
memory_custom_instance.llm.generate_response.return_value = '{"facts": ["fact1", "fact2"]}'
memory_custom_instance.embedding_model = MockEmbeddings()
with patch("mem0.memory.main.parse_messages", return_value="Test message") as mock_parse_messages:
with patch(

7
tests/vector_stores/test_baidu.py
View File

@@ -2,10 +2,9 @@ from unittest.mock import Mock, patch, PropertyMock
import pytest
from mem0.vector_stores.baidu import BaiduDB, OutputData
from pymochow.model.enum import MetricType, TableState, ServerErrCode
from pymochow.model.schema import Field, Schema, VectorIndex, FilteringIndex, HNSWParams, AutoBuildRowCountIncrement
from pymochow.model.table import Partition, Row, VectorSearchConfig, VectorTopkSearchRequest, FloatVector, Table
from mem0.vector_stores.baidu import BaiduDB
from pymochow.model.enum import TableState, ServerErrCode
from pymochow.model.table import VectorSearchConfig, VectorTopkSearchRequest, FloatVector, Table
from pymochow.exception import ServerError

48
tests/vector_stores/test_mongodb.py
View File

@@ -1,7 +1,7 @@
import pytest
from unittest.mock import MagicMock, patch
from mem0.vector_stores.mongodb import MongoDB
from pymongo.operations import SearchIndexModel
@pytest.fixture
@patch("mem0.vector_stores.mongodb.MongoClient")
@@ -19,10 +19,11 @@ def mongo_vector_fixture(mock_mongo_client):
db_name="test_db",
collection_name="test_collection",
embedding_model_dims=1536,
mongo_uri="mongodb://username:password@localhost:27017"
mongo_uri="mongodb://username:password@localhost:27017",
)
return mongo_vector, mock_collection, mock_db
def test_initalize_create_col(mongo_vector_fixture):
mongo_vector, mock_collection, mock_db = mongo_vector_fixture
assert mongo_vector.collection_name == "test_collection"
@@ -49,12 +50,13 @@ def test_initalize_create_col(mongo_vector_fixture):
"dimensions": 1536,
"similarity": "cosine",
}
}
},
}
}
},
}
assert mongo_vector.collection == mock_collection
def test_insert(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
vectors = [[0.1] * 1536, [0.2] * 1536]
@@ -62,12 +64,13 @@ def test_insert(mongo_vector_fixture):
ids = ["id1", "id2"]
mongo_vector.insert(vectors, payloads, ids)
expected_records=[
expected_records = [
({"_id": ids[0], "embedding": vectors[0], "payload": payloads[0]}),
({"_id": ids[1], "embedding": vectors[1], "payload": payloads[1]})
({"_id": ids[1], "embedding": vectors[1], "payload": payloads[1]}),
]
mock_collection.insert_many.assert_called_once_with(expected_records)
def test_search(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
query_vector = [0.1] * 1536
@@ -79,25 +82,28 @@ def test_search(mongo_vector_fixture):
results = mongo_vector.search("query_str", query_vector, limit=2)
mock_collection.list_search_indexes.assert_called_with(name="test_collection_vector_index")
mock_collection.aggregate.assert_called_once_with([
{
"$vectorSearch": {
"index": "test_collection_vector_index",
"limit": 2,
"numCandidates": 2,
"queryVector": query_vector,
"path": "embedding",
mock_collection.aggregate.assert_called_once_with(
[
{
"$vectorSearch": {
"index": "test_collection_vector_index",
"limit": 2,
"numCandidates": 2,
"queryVector": query_vector,
"path": "embedding",
},
},
},
{"$set": {"score": {"$meta": "vectorSearchScore"}}},
{"$project": {"embedding": 0}},
])
{"$set": {"score": {"$meta": "vectorSearchScore"}}},
{"$project": {"embedding": 0}},
]
)
assert len(results) == 2
assert results[0].id == "id1"
assert results[0].score == 0.9
assert results[1].id == "id2"
assert results[1].score == 0.8
def test_delete(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
mock_delete_result = MagicMock()
@@ -107,6 +113,7 @@ def test_delete(mongo_vector_fixture):
mongo_vector.delete("id1")
mock_collection.delete_one.assert_called_with({"_id": "id1"})
def test_update(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
mock_update_result = MagicMock()
@@ -122,6 +129,7 @@ def test_update(mongo_vector_fixture):
{"$set": {"embedding": vectorValue, "payload": payloadValue}},
)
def test_get(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
mock_collection.find_one.return_value = {"_id": "id1", "payload": {"key": "value1"}}
@@ -131,6 +139,7 @@ def test_get(mongo_vector_fixture):
assert result.id == "id1"
assert result.payload == {"key": "value1"}
def test_list_cols(mongo_vector_fixture):
mongo_vector, _, mock_db = mongo_vector_fixture
mock_db.list_collection_names.return_value = ["col1", "col2"]
@@ -138,12 +147,14 @@ def test_list_cols(mongo_vector_fixture):
collections = mongo_vector.list_cols()
assert collections == ["col1", "col2"]
def test_delete_col(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
mongo_vector.delete_col()
mock_collection.drop.assert_called_once()
def test_col_info(mongo_vector_fixture):
mongo_vector, _, mock_db = mongo_vector_fixture
mock_db.command.return_value = {"count": 10, "size": 1024}
@@ -154,6 +165,7 @@ def test_col_info(mongo_vector_fixture):
assert info["count"] == 10
assert info["size"] == 1024
def test_list(mongo_vector_fixture):
mongo_vector, mock_collection, _ = mongo_vector_fixture
mock_cursor = MagicMock()

2
tests/vector_stores/test_pinecone.py
View File

@@ -88,7 +88,7 @@ def test_get_vector_found(pinecone_db):
# or a list of dictionaries, not a dictionary with an 'id' field
# Create a mock Vector object
from pinecone.data.dataclasses.vector import Vector
from pinecone import Vector
mock_vector = Vector(id="id1", values=[0.1] * 128, metadata={"name": "vector1"})