Code formatting (#1986)

tests/embeddings/test_azure_openai_embeddings.py

@@ -33,11 +33,7 @@ def test_embed_text(mock_openai_client):
 
 @pytest.mark.parametrize(
     "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):
     config = BaseEmbedderConfig(
@@ -47,8 +43,8 @@ def test_embed_text_with_default_headers(default_headers, expected_header):
             "api_version": "test_version",
             "azure_endpoint": "test_endpoint",
             "azuer_deployment": "test_deployment",
-            "default_headers": default_headers
-        }
+            "default_headers": default_headers,
+        },
     )
     embedder = AzureOpenAIEmbedding(config)
     assert embedder.client.api_key == "test"

tests/embeddings/test_gemini.py

@@ -12,17 +12,11 @@ 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")
 
 
 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
 
     embedder = GoogleGenAIEmbedding(config)
@@ -31,7 +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!")

tests/embeddings/test_vertexai_embeddings.py

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