my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	from unittest.mock import Mock, patch
#2
#3	import numpy as np
#4	import pytest
#5
#6	from mem0.configs.embeddings.base import BaseEmbedderConfig
#7	from mem0.embeddings.huggingface import HuggingFaceEmbedding
#8
#9
#10	@pytest.fixture
#11	def mock_sentence_transformer():
#12	with patch("mem0.embeddings.huggingface.SentenceTransformer") as mock_transformer:
#13	mock_model = Mock()
#14	mock_transformer.return_value = mock_model
#15	yield mock_model
#16
#17
#18	def test_embed_default_model(mock_sentence_transformer):
#19	config = BaseEmbedderConfig()
#20	embedder = HuggingFaceEmbedding(config)
#21
#22	mock_sentence_transformer.encode.return_value = np.array([0.1, 0.2, 0.3])
#23	result = embedder.embed("Hello world")
#24
#25	mock_sentence_transformer.encode.assert_called_once_with("Hello world", convert_to_numpy=True)
#26	assert result == [0.1, 0.2, 0.3]
#27
#28
#29	def test_embed_custom_model(mock_sentence_transformer):
#30	config = BaseEmbedderConfig(model="paraphrase-MiniLM-L6-v2")
#31	embedder = HuggingFaceEmbedding(config)
#32
#33	mock_sentence_transformer.encode.return_value = np.array([0.4, 0.5, 0.6])
#34	result = embedder.embed("Custom model test")
#35
#36	mock_sentence_transformer.encode.assert_called_once_with("Custom model test", convert_to_numpy=True)
#37	assert result == [0.4, 0.5, 0.6]
#38
#39
#40	def test_embed_with_model_kwargs(mock_sentence_transformer):
#41	config = BaseEmbedderConfig(model="all-MiniLM-L6-v2", model_kwargs={"device": "cuda"})
#42	embedder = HuggingFaceEmbedding(config)
#43
#44	mock_sentence_transformer.encode.return_value = np.array([0.7, 0.8, 0.9])
#45	result = embedder.embed("Test with device")
#46
#47	mock_sentence_transformer.encode.assert_called_once_with("Test with device", convert_to_numpy=True)
#48	assert result == [0.7, 0.8, 0.9]
#49
#50
#51	def test_embed_sets_embedding_dims(mock_sentence_transformer):
#52	config = BaseEmbedderConfig()
#53
#54	mock_sentence_transformer.get_sentence_embedding_dimension.return_value = 384
#55	embedder = HuggingFaceEmbedding(config)
#56
#57	assert embedder.config.embedding_dims == 384
#58	mock_sentence_transformer.get_sentence_embedding_dimension.assert_called_once()
#59
#60
#61	def test_embed_with_custom_embedding_dims(mock_sentence_transformer):
#62	config = BaseEmbedderConfig(model="all-mpnet-base-v2", embedding_dims=768)
#63	embedder = HuggingFaceEmbedding(config)
#64
#65	mock_sentence_transformer.encode.return_value = np.array([1.0, 1.1, 1.2])
#66	result = embedder.embed("Custom embedding dims")
#67
#68	mock_sentence_transformer.encode.assert_called_once_with("Custom embedding dims", convert_to_numpy=True)
#69
#70	assert embedder.config.embedding_dims == 768
#71
#72	assert result == [1.0, 1.1, 1.2]
#73
#74
#75	def test_embed_with_huggingface_base_url():
#76	config = BaseEmbedderConfig(
#77	huggingface_base_url="http://localhost:8080",
#78	model="my-custom-model",
#79	model_kwargs={"truncate": True},
#80	)
#81	with patch("mem0.embeddings.huggingface.OpenAI") as mock_openai:
#82	mock_client = Mock()
#83	mock_openai.return_value = mock_client
#84
#85	# Create a mock for the response object and its attributes
#86	mock_embedding_response = Mock()
#87	mock_embedding_response.embedding = [0.1, 0.2, 0.3]
#88
#89	mock_create_response = Mock()
#90	mock_create_response.data = [mock_embedding_response]
#91
#92	mock_client.embeddings.create.return_value = mock_create_response
#93
#94	embedder = HuggingFaceEmbedding(config)
#95	result = embedder.embed("Hello from custom endpoint")
#96
#97	mock_openai.assert_called_once_with(base_url="http://localhost:8080")
#98	mock_client.embeddings.create.assert_called_once_with(
#99	input="Hello from custom endpoint",
#100	model="my-custom-model",
#101	truncate=True,
#102	)
#103	assert result == [0.1, 0.2, 0.3]
#104

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public