From 8c6d16a6f011d42d3391a751b93b0083db06f4ff Mon Sep 17 00:00:00 2001 From: "Rafael Nico T. Maniquiz" Date: Wed, 5 Mar 2025 22:53:36 +0800 Subject: [PATCH] Fix Embedding Dimension Parameter Not Being Passed (#2304) --- mem0/embeddings/gemini.py | 2 +- mem0/embeddings/openai.py | 2 +- tests/embeddings/test_gemini.py | 4 ++-- tests/embeddings/test_openai_embeddings.py | 10 +++++----- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/mem0/embeddings/gemini.py b/mem0/embeddings/gemini.py index d549db78..a042af35 100644 --- a/mem0/embeddings/gemini.py +++ b/mem0/embeddings/gemini.py @@ -28,5 +28,5 @@ class GoogleGenAIEmbedding(EmbeddingBase): list: The embedding vector. """ text = text.replace("\n", " ") - response = genai.embed_content(model=self.config.model, content=text) + response = genai.embed_content(model=self.config.model, content=text, output_dimensionality=self.config.embedding_dims) return response["embedding"] diff --git a/mem0/embeddings/openai.py b/mem0/embeddings/openai.py index a774544c..649f257b 100644 --- a/mem0/embeddings/openai.py +++ b/mem0/embeddings/openai.py @@ -29,4 +29,4 @@ class OpenAIEmbedding(EmbeddingBase): list: The embedding vector. """ text = text.replace("\n", " ") - return self.client.embeddings.create(input=[text], model=self.config.model).data[0].embedding + return self.client.embeddings.create(input=[text], model=self.config.model, dimensions = self.config.embedding_dims).data[0].embedding diff --git a/tests/embeddings/test_gemini.py b/tests/embeddings/test_gemini.py index a1ae50b1..b9fdc8c5 100644 --- a/tests/embeddings/test_gemini.py +++ b/tests/embeddings/test_gemini.py @@ -12,7 +12,7 @@ def mock_genai(): @pytest.fixture def config(): - return BaseEmbedderConfig(api_key="dummy_api_key", model="test_model") + return BaseEmbedderConfig(api_key="dummy_api_key", model="test_model", embedding_dims=786) def test_embed_query(mock_genai, config): @@ -25,4 +25,4 @@ 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!") + mock_genai.assert_called_once_with(model="test_model", content="Hello, world!", output_dimensionality=786) diff --git a/tests/embeddings/test_openai_embeddings.py b/tests/embeddings/test_openai_embeddings.py index 113a8e64..bd5a11df 100644 --- a/tests/embeddings/test_openai_embeddings.py +++ b/tests/embeddings/test_openai_embeddings.py @@ -21,7 +21,7 @@ def test_embed_default_model(mock_openai_client): result = embedder.embed("Hello world") - mock_openai_client.embeddings.create.assert_called_once_with(input=["Hello world"], model="text-embedding-3-small") + mock_openai_client.embeddings.create.assert_called_once_with(input=["Hello world"], model="text-embedding-3-small", dimensions = 1536) assert result == [0.1, 0.2, 0.3] @@ -35,7 +35,7 @@ def test_embed_custom_model(mock_openai_client): result = embedder.embed("Test embedding") mock_openai_client.embeddings.create.assert_called_once_with( - input=["Test embedding"], model="text-embedding-2-medium" + input=["Test embedding"], model="text-embedding-2-medium", dimensions = 1024 ) assert result == [0.4, 0.5, 0.6] @@ -49,7 +49,7 @@ def test_embed_removes_newlines(mock_openai_client): result = embedder.embed("Hello\nworld") - mock_openai_client.embeddings.create.assert_called_once_with(input=["Hello world"], model="text-embedding-3-small") + mock_openai_client.embeddings.create.assert_called_once_with(input=["Hello world"], model="text-embedding-3-small", dimensions = 1536) assert result == [0.7, 0.8, 0.9] @@ -63,7 +63,7 @@ def test_embed_without_api_key_env_var(mock_openai_client): result = embedder.embed("Testing API key") mock_openai_client.embeddings.create.assert_called_once_with( - input=["Testing API key"], model="text-embedding-3-small" + input=["Testing API key"], model="text-embedding-3-small", dimensions = 1536 ) assert result == [1.0, 1.1, 1.2] @@ -79,6 +79,6 @@ def test_embed_uses_environment_api_key(mock_openai_client, monkeypatch): result = embedder.embed("Environment key test") mock_openai_client.embeddings.create.assert_called_once_with( - input=["Environment key test"], model="text-embedding-3-small" + input=["Environment key test"], model="text-embedding-3-small", dimensions = 1536 ) assert result == [1.3, 1.4, 1.5]