my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	from unittest.mock import Mock, patch
#2
#3	import pytest
#4
#5	from mem0.embeddings.vertexai import VertexAIEmbedding
#6
#7
#8	@pytest.fixture
#9	def mock_text_embedding_model():
#10	with patch("mem0.embeddings.vertexai.TextEmbeddingModel") as mock_model:
#11	mock_instance = Mock()
#12	mock_model.from_pretrained.return_value = mock_instance
#13	yield mock_instance
#14
#15
#16	@pytest.fixture
#17	def mock_os_environ():
#18	with patch("mem0.embeddings.vertexai.os.environ", {}) as mock_environ:
#19	yield mock_environ
#20
#21
#22	@pytest.fixture
#23	def mock_config():
#24	with patch("mem0.configs.embeddings.base.BaseEmbedderConfig") as mock_config:
#25	mock_config.return_value.vertex_credentials_json = "/path/to/credentials.json"
#26	yield mock_config
#27
#28
#29	@pytest.fixture
#30	def mock_embedding_types():
#31	return [
#32	"SEMANTIC_SIMILARITY",
#33	"CLASSIFICATION",
#34	"CLUSTERING",
#35	"RETRIEVAL_DOCUMENT",
#36	"RETRIEVAL_QUERY",
#37	"QUESTION_ANSWERING",
#38	"FACT_VERIFICATION",
#39	"CODE_RETRIEVAL_QUERY",
#40	]
#41
#42
#43	@pytest.fixture
#44	def mock_text_embedding_input():
#45	with patch("mem0.embeddings.vertexai.TextEmbeddingInput") as mock_input:
#46	yield mock_input
#47
#48
#49	@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
#50	def test_embed_default_model(mock_text_embedding_model, mock_os_environ, mock_config, mock_text_embedding_input):
#51	mock_config.return_value.model = "text-embedding-004"
#52	mock_config.return_value.embedding_dims = 256
#53
#54	config = mock_config()
#55	embedder = VertexAIEmbedding(config)
#56
#57	mock_embedding = Mock(values=[0.1, 0.2, 0.3])
#58	mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
#59
#60	embedder.embed("Hello world")
#61	mock_text_embedding_input.assert_called_once_with(text="Hello world", task_type="SEMANTIC_SIMILARITY")
#62	mock_text_embedding_model.from_pretrained.assert_called_once_with("text-embedding-004")
#63
#64	mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with(
#65	texts=[mock_text_embedding_input("Hello world")], output_dimensionality=256
#66	)
#67
#68
#69	@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
#70	def test_embed_custom_model(mock_text_embedding_model, mock_os_environ, mock_config, mock_text_embedding_input):
#71	mock_config.return_value.model = "custom-embedding-model"
#72	mock_config.return_value.embedding_dims = 512
#73
#74	config = mock_config()
#75
#76	embedder = VertexAIEmbedding(config)
#77
#78	mock_embedding = Mock(values=[0.4, 0.5, 0.6])
#79	mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
#80
#81	result = embedder.embed("Test embedding")
#82	mock_text_embedding_input.assert_called_once_with(text="Test embedding", task_type="SEMANTIC_SIMILARITY")
#83	mock_text_embedding_model.from_pretrained.assert_called_with("custom-embedding-model")
#84	mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_once_with(
#85	texts=[mock_text_embedding_input("Test embedding")], output_dimensionality=512
#86	)
#87
#88	assert result == [0.4, 0.5, 0.6]
#89
#90
#91	@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
#92	def test_embed_with_memory_action(
#93	mock_text_embedding_model, mock_os_environ, mock_config, mock_embedding_types, mock_text_embedding_input
#94	):
#95	mock_config.return_value.model = "text-embedding-004"
#96	mock_config.return_value.embedding_dims = 256
#97
#98	for embedding_type in mock_embedding_types:
#99	mock_config.return_value.memory_add_embedding_type = embedding_type
#100	mock_config.return_value.memory_update_embedding_type = embedding_type
#101	mock_config.return_value.memory_search_embedding_type = embedding_type
#102
#103	config = mock_config()
#104	embedder = VertexAIEmbedding(config)
#105
#106	mock_text_embedding_model.from_pretrained.assert_called_with("text-embedding-004")
#107
#108	for memory_action in ["add", "update", "search"]:
#109	embedder.embed("Hello world", memory_action=memory_action)
#110
#111	mock_text_embedding_input.assert_called_with(text="Hello world", task_type=embedding_type)
#112	mock_text_embedding_model.from_pretrained.return_value.get_embeddings.assert_called_with(
#113	texts=[mock_text_embedding_input("Hello world", embedding_type)], output_dimensionality=256
#114	)
#115
#116
#117	@patch("mem0.embeddings.vertexai.os")
#118	def test_credentials_from_environment(mock_os, mock_text_embedding_model, mock_config):
#119	mock_config.vertex_credentials_json = None
#120	config = mock_config()
#121	VertexAIEmbedding(config)
#122
#123	mock_os.environ.setitem.assert_not_called()
#124
#125
#126	@patch("mem0.embeddings.vertexai.os")
#127	def test_missing_credentials(mock_os, mock_text_embedding_model, mock_config):
#128	mock_os.getenv.return_value = None
#129	mock_config.return_value.vertex_credentials_json = None
#130
#131	config = mock_config()
#132
#133	with pytest.raises(ValueError, match="Google application credentials JSON is not provided"):
#134	VertexAIEmbedding(config)
#135
#136
#137	@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
#138	def test_embed_with_different_dimensions(mock_text_embedding_model, mock_os_environ, mock_config):
#139	mock_config.return_value.embedding_dims = 1024
#140
#141	config = mock_config()
#142	embedder = VertexAIEmbedding(config)
#143
#144	mock_embedding = Mock(values=[0.1] * 1024)
#145	mock_text_embedding_model.from_pretrained.return_value.get_embeddings.return_value = [mock_embedding]
#146
#147	result = embedder.embed("Large embedding test")
#148
#149	assert result == [0.1] * 1024
#150
#151
#152	@patch("mem0.embeddings.vertexai.TextEmbeddingModel")
#153	def test_invalid_memory_action(mock_text_embedding_model, mock_config):
#154	mock_config.return_value.model = "text-embedding-004"
#155	mock_config.return_value.embedding_dims = 256
#156
#157	config = mock_config()
#158	embedder = VertexAIEmbedding(config)
#159
#160	with pytest.raises(ValueError):
#161	embedder.embed("Hello world", memory_action="invalid_action")
#162

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public