Formatting (#2750)

tests/vector_stores/test_azure_ai_search.py

@@ -13,13 +13,15 @@ from mem0.vector_stores.azure_ai_search import AzureAISearch
 # Fixture to patch SearchClient and SearchIndexClient and create an instance of AzureAISearch.
 @pytest.fixture
 def mock_clients():
-    with patch("mem0.vector_stores.azure_ai_search.SearchClient") as MockSearchClient, \
-         patch("mem0.vector_stores.azure_ai_search.SearchIndexClient") as MockIndexClient, \
-         patch("mem0.vector_stores.azure_ai_search.AzureKeyCredential") as MockAzureKeyCredential:
+    with (
+        patch("mem0.vector_stores.azure_ai_search.SearchClient") as MockSearchClient,
+        patch("mem0.vector_stores.azure_ai_search.SearchIndexClient") as MockIndexClient,
+        patch("mem0.vector_stores.azure_ai_search.AzureKeyCredential") as MockAzureKeyCredential,
+    ):
         # Create mocked instances for search and index clients.
         mock_search_client = MockSearchClient.return_value
         mock_index_client = MockIndexClient.return_value
-        
+
         # Mock the client._client._config.user_agent_policy.add_user_agent
         mock_search_client._client = MagicMock()
         mock_search_client._client._config.user_agent_policy.add_user_agent = Mock()
@@ -62,7 +64,7 @@ def azure_ai_search_instance(mock_clients):
         api_key="test-api-key",
         embedding_model_dims=3,
         compression_type="binary",  # testing binary quantization option
-        use_float16=True
+        use_float16=True,
     )
     # Return instance and clients for verification.
     return instance, mock_search_client, mock_index_client
@@ -70,21 +72,18 @@ def azure_ai_search_instance(mock_clients):
 
 # --- Tests for AzureAISearchConfig ---
 
+
 def test_config_validation_valid():
     """Test valid configurations are accepted."""
     # Test minimal configuration
-    config = AzureAISearchConfig(
-        service_name="test-service",
-        api_key="test-api-key",
-        embedding_model_dims=768
-    )
+    config = AzureAISearchConfig(service_name="test-service", api_key="test-api-key", embedding_model_dims=768)
     assert config.collection_name == "mem0"  # Default value
     assert config.service_name == "test-service"
     assert config.api_key == "test-api-key"
     assert config.embedding_model_dims == 768
     assert config.compression_type is None
     assert config.use_float16 is False
-    
+
     # Test with all optional parameters
     config = AzureAISearchConfig(
         collection_name="custom-index",
@@ -92,7 +91,7 @@ def test_config_validation_valid():
         api_key="test-api-key",
         embedding_model_dims=1536,
         compression_type="scalar",
-        use_float16=True
+        use_float16=True,
     )
     assert config.collection_name == "custom-index"
     assert config.compression_type == "scalar"
@@ -106,7 +105,7 @@ def test_config_validation_invalid_compression_type():
             service_name="test-service",
             api_key="test-api-key",
             embedding_model_dims=768,
-            compression_type="invalid-type"  # Not a valid option
+            compression_type="invalid-type",  # Not a valid option
         )
     assert "Invalid compression_type" in str(exc_info.value)
 
@@ -118,7 +117,7 @@ def test_config_validation_deprecated_use_compression():
             service_name="test-service",
             api_key="test-api-key",
             embedding_model_dims=768,
-            use_compression=True  # Deprecated parameter
+            use_compression=True,  # Deprecated parameter
         )
     # Fix: Use a partial string match instead of exact match
     assert "use_compression" in str(exc_info.value)
@@ -132,7 +131,7 @@ def test_config_validation_extra_fields():
             service_name="test-service",
             api_key="test-api-key",
             embedding_model_dims=768,
-            unknown_parameter="value"  # Extra field
+            unknown_parameter="value",  # Extra field
         )
     assert "Extra fields not allowed" in str(exc_info.value)
     assert "unknown_parameter" in str(exc_info.value)
@@ -140,30 +139,28 @@ def test_config_validation_extra_fields():
 
 # --- Tests for AzureAISearch initialization ---
 
+
 def test_initialization(mock_clients):
     """Test AzureAISearch initialization with different parameters."""
     mock_search_client, mock_index_client, mock_azure_key_credential = mock_clients
-    
+
     # Test with minimal parameters
     instance = AzureAISearch(
-        service_name="test-service",
-        collection_name="test-index",
-        api_key="test-api-key",
-        embedding_model_dims=768
+        service_name="test-service", collection_name="test-index", api_key="test-api-key", embedding_model_dims=768
     )
-    
+
     # Verify initialization parameters
     assert instance.index_name == "test-index"
     assert instance.collection_name == "test-index"
     assert instance.embedding_model_dims == 768
     assert instance.compression_type == "none"  # Default when None is passed
     assert instance.use_float16 is False
-    
+
     # Verify client creation
     mock_azure_key_credential.assert_called_with("test-api-key")
     assert "mem0" in mock_search_client._client._config.user_agent_policy.add_user_agent.call_args[0]
     assert "mem0" in mock_index_client._client._config.user_agent_policy.add_user_agent.call_args[0]
-    
+
     # Verify index creation was called
     mock_index_client.create_or_update_index.assert_called_once()
 
@@ -171,75 +168,75 @@ def test_initialization(mock_clients):
 def test_initialization_with_compression_types(mock_clients):
     """Test initialization with different compression types."""
     mock_search_client, mock_index_client, _ = mock_clients
-    
+
     # Test with scalar compression
     instance = AzureAISearch(
         service_name="test-service",
         collection_name="scalar-index",
         api_key="test-api-key",
         embedding_model_dims=768,
-        compression_type="scalar"
+        compression_type="scalar",
     )
     assert instance.compression_type == "scalar"
-    
+
     # Capture the index creation call
     args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
     index = args[0]
     # Verify scalar compression was configured
-    assert hasattr(index.vector_search, 'compressions')
+    assert hasattr(index.vector_search, "compressions")
     assert len(index.vector_search.compressions) > 0
     assert "ScalarQuantizationCompression" in str(type(index.vector_search.compressions[0]))
-    
+
     # Test with binary compression
     instance = AzureAISearch(
         service_name="test-service",
         collection_name="binary-index",
         api_key="test-api-key",
         embedding_model_dims=768,
-        compression_type="binary"
+        compression_type="binary",
     )
     assert instance.compression_type == "binary"
-    
+
     # Capture the index creation call
     args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
     index = args[0]
     # Verify binary compression was configured
-    assert hasattr(index.vector_search, 'compressions')
+    assert hasattr(index.vector_search, "compressions")
     assert len(index.vector_search.compressions) > 0
     assert "BinaryQuantizationCompression" in str(type(index.vector_search.compressions[0]))
-    
+
     # Test with no compression
     instance = AzureAISearch(
         service_name="test-service",
         collection_name="no-compression-index",
         api_key="test-api-key",
         embedding_model_dims=768,
-        compression_type=None
+        compression_type=None,
     )
     assert instance.compression_type == "none"
-    
+
     # Capture the index creation call
     args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
     index = args[0]
     # Verify no compression was configured
-    assert hasattr(index.vector_search, 'compressions')
+    assert hasattr(index.vector_search, "compressions")
     assert len(index.vector_search.compressions) == 0
 
 
 def test_initialization_with_float_precision(mock_clients):
     """Test initialization with different float precision settings."""
     mock_search_client, mock_index_client, _ = mock_clients
-    
+
     # Test with half precision (float16)
     instance = AzureAISearch(
         service_name="test-service",
         collection_name="float16-index",
         api_key="test-api-key",
         embedding_model_dims=768,
-        use_float16=True
+        use_float16=True,
     )
     assert instance.use_float16 is True
-    
+
     # Capture the index creation call
     args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
     index = args[0]
@@ -247,17 +244,17 @@ def test_initialization_with_float_precision(mock_clients):
     vector_field = next((f for f in index.fields if f.name == "vector"), None)
     assert vector_field is not None
     assert "Edm.Half" in vector_field.type
-    
+
     # Test with full precision (float32)
     instance = AzureAISearch(
         service_name="test-service",
         collection_name="float32-index",
         api_key="test-api-key",
         embedding_model_dims=768,
-        use_float16=False
+        use_float16=False,
     )
     assert instance.use_float16 is False
-    
+
     # Capture the index creation call
     args, _ = mock_index_client.create_or_update_index.call_args_list[-1]
     index = args[0]
@@ -269,21 +266,22 @@ def test_initialization_with_float_precision(mock_clients):
 
 # --- Tests for create_col method ---
 
+
 def test_create_col(azure_ai_search_instance):
     """Test the create_col method creates an index with the correct configuration."""
     instance, _, mock_index_client = azure_ai_search_instance
-    
+
     # create_col is called during initialization, so we check the call that was already made
     mock_index_client.create_or_update_index.assert_called_once()
-    
+
     # Verify the index configuration
     args, _ = mock_index_client.create_or_update_index.call_args
     index = args[0]
-    
+
     # Check basic properties
     assert index.name == "test-index"
     assert len(index.fields) == 6  # id, user_id, run_id, agent_id, vector, payload
-    
+
     # Check that required fields are present
     field_names = [f.name for f in index.fields]
     assert "id" in field_names
@@ -292,22 +290,22 @@ def test_create_col(azure_ai_search_instance):
     assert "user_id" in field_names
     assert "run_id" in field_names
     assert "agent_id" in field_names
-    
+
     # Check that id is the key field
     id_field = next(f for f in index.fields if f.name == "id")
     assert id_field.key is True
-    
+
     # Check vector search configuration
     assert index.vector_search is not None
     assert len(index.vector_search.profiles) == 1
     assert index.vector_search.profiles[0].name == "my-vector-config"
     assert index.vector_search.profiles[0].algorithm_configuration_name == "my-algorithms-config"
-    
+
     # Check algorithms
     assert len(index.vector_search.algorithms) == 1
     assert index.vector_search.algorithms[0].name == "my-algorithms-config"
     assert "HnswAlgorithmConfiguration" in str(type(index.vector_search.algorithms[0]))
-    
+
     # With binary compression and float16, we should have compression configuration
     assert len(index.vector_search.compressions) == 1
     assert index.vector_search.compressions[0].compression_name == "myCompression"
@@ -317,24 +315,24 @@ def test_create_col(azure_ai_search_instance):
 def test_create_col_scalar_compression(mock_clients):
     """Test creating a collection with scalar compression."""
     mock_search_client, mock_index_client, _ = mock_clients
-    
+
     AzureAISearch(
         service_name="test-service",
         collection_name="scalar-index",
         api_key="test-api-key",
         embedding_model_dims=768,
-        compression_type="scalar"
+        compression_type="scalar",
     )
-    
+
     # Verify the index configuration
     args, _ = mock_index_client.create_or_update_index.call_args
     index = args[0]
-    
+
     # Check compression configuration
     assert len(index.vector_search.compressions) == 1
     assert index.vector_search.compressions[0].compression_name == "myCompression"
     assert "ScalarQuantizationCompression" in str(type(index.vector_search.compressions[0]))
-    
+
     # Check profile references compression
     assert index.vector_search.profiles[0].compression_name == "myCompression"
 
@@ -342,28 +340,29 @@ def test_create_col_scalar_compression(mock_clients):
 def test_create_col_no_compression(mock_clients):
     """Test creating a collection with no compression."""
     mock_search_client, mock_index_client, _ = mock_clients
-    
+
     AzureAISearch(
         service_name="test-service",
         collection_name="no-compression-index",
         api_key="test-api-key",
         embedding_model_dims=768,
-        compression_type=None
+        compression_type=None,
     )
-    
+
     # Verify the index configuration
     args, _ = mock_index_client.create_or_update_index.call_args
     index = args[0]
-    
+
     # Check compression configuration - should be empty
     assert len(index.vector_search.compressions) == 0
-    
+
     # Check profile doesn't reference compression
     assert index.vector_search.profiles[0].compression_name is None
 
 
 # --- Tests for insert method ---
 
+
 def test_insert_single(azure_ai_search_instance):
     """Test inserting a single vector."""
     instance, mock_search_client, _ = azure_ai_search_instance
@@ -372,9 +371,7 @@ def test_insert_single(azure_ai_search_instance):
     ids = ["doc1"]
 
     # Fix: Include status_code: 201 in mock response
-    mock_search_client.upload_documents.return_value = [
-        {"status": True, "id": "doc1", "status_code": 201}
-    ]
+    mock_search_client.upload_documents.return_value = [{"status": True, "id": "doc1", "status_code": 201}]
 
     instance.insert(vectors, payloads, ids)
 
@@ -396,35 +393,35 @@ def test_insert_single(azure_ai_search_instance):
 def test_insert_multiple(azure_ai_search_instance):
     """Test inserting multiple vectors in one call."""
     instance, mock_search_client, _ = azure_ai_search_instance
-    
+
     # Create multiple vectors
     num_docs = 3
-    vectors = [[float(i)/10, float(i+1)/10, float(i+2)/10] for i in range(num_docs)]
+    vectors = [[float(i) / 10, float(i + 1) / 10, float(i + 2) / 10] for i in range(num_docs)]
     payloads = [{"user_id": f"user{i}", "content": f"Test content {i}"} for i in range(num_docs)]
     ids = [f"doc{i}" for i in range(num_docs)]
-    
+
     # Configure mock to return success for all documents (fix: add status_code 201)
     mock_search_client.upload_documents.return_value = [
         {"status": True, "id": id_val, "status_code": 201} for id_val in ids
     ]
-    
+
     # Insert the documents
     instance.insert(vectors, payloads, ids)
-    
+
     # Verify upload_documents was called with correct documents
     mock_search_client.upload_documents.assert_called_once()
     args, _ = mock_search_client.upload_documents.call_args
     documents = args[0]
-    
+
     # Verify all documents were included
     assert len(documents) == num_docs
-    
+
     # Check first document
     assert documents[0]["id"] == "doc0"
     assert documents[0]["vector"] == [0.0, 0.1, 0.2]
     assert documents[0]["payload"] == json.dumps(payloads[0])
     assert documents[0]["user_id"] == "user0"
-    
+
     # Check last document
     assert documents[2]["id"] == "doc2"
     assert documents[2]["vector"] == [0.2, 0.3, 0.4]
@@ -437,9 +434,7 @@ def test_insert_with_error(azure_ai_search_instance):
     instance, mock_search_client, _ = azure_ai_search_instance
 
     # Configure mock to return an error for one document
-    mock_search_client.upload_documents.return_value = [
-        {"status": False, "id": "doc1", "errorMessage": "Azure error"}
-    ]
+    mock_search_client.upload_documents.return_value = [{"status": False, "id": "doc1", "errorMessage": "Azure error"}]
 
     vectors = [[0.1, 0.2, 0.3]]
     payloads = [{"user_id": "user1"}]
@@ -454,7 +449,7 @@ def test_insert_with_error(azure_ai_search_instance):
     # Configure mock to return mixed success/failure for multiple documents
     mock_search_client.upload_documents.return_value = [
         {"status": True, "id": "doc1"},  # This should not cause failure
-        {"status": False, "id": "doc2", "errorMessage": "Azure error"}
+        {"status": False, "id": "doc2", "errorMessage": "Azure error"},
     ]
 
     vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
@@ -465,8 +460,9 @@ def test_insert_with_error(azure_ai_search_instance):
     with pytest.raises(Exception) as exc_info:
         instance.insert(vectors, payloads, ids)
 
-    assert "Insert failed for document doc2" in str(exc_info.value) or \
-           "Insert failed for document doc1" in str(exc_info.value)
+    assert "Insert failed for document doc2" in str(exc_info.value) or "Insert failed for document doc1" in str(
+        exc_info.value
+    )
 
 
 def test_insert_with_missing_payload_fields(azure_ai_search_instance):
@@ -500,23 +496,24 @@ def test_insert_with_missing_payload_fields(azure_ai_search_instance):
 def test_insert_with_http_error(azure_ai_search_instance):
     """Test insert when Azure client throws an HTTP error."""
     instance, mock_search_client, _ = azure_ai_search_instance
-    
+
     # Configure mock to raise an HttpResponseError
     mock_search_client.upload_documents.side_effect = HttpResponseError("Azure service error")
-    
+
     vectors = [[0.1, 0.2, 0.3]]
     payloads = [{"user_id": "user1"}]
     ids = ["doc1"]
-    
+
     # Insert should propagate the HTTP error
     with pytest.raises(HttpResponseError) as exc_info:
         instance.insert(vectors, payloads, ids)
-    
+
     assert "Azure service error" in str(exc_info.value)
 
 
 # --- Tests for search method ---
 
+
 def test_search_basic(azure_ai_search_instance):
     """Test basic vector search without filters."""
     instance, mock_search_client, _ = azure_ai_search_instance
@@ -536,9 +533,7 @@ def test_search_basic(azure_ai_search_instance):
     # Search with a vector
     query_text = "test query"  # Add a query string
     query_vector = [0.1, 0.2, 0.3]
-    results = instance.search(
-        query_text, query_vector, limit=5
-    )  # Pass the query string
+    results = instance.search(query_text, query_vector, limit=5)  # Pass the query string
 
     # Verify search was called correctly
     mock_search_client.search.assert_called_once()