From 295d8b99c369d8a45ae7f2cd54415cd9e69748b1 Mon Sep 17 00:00:00 2001 From: Mustafa Elbehery Date: Tue, 9 Sep 2025 22:20:02 +0200 Subject: [PATCH] refactor(client): replace all AsyncMilvusClient usage of has_collection() with list_collections() Signed-off-by: Mustafa Elbehery --- .../remote/vector_io/milvus/milvus.py | 9 ++++++--- .../providers/vector_io/remote/test_milvus.py | 18 +++++++++--------- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index 383fa517a..2e9ed9ced 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -67,7 +67,8 @@ class MilvusIndex(EmbeddingIndex): async def delete(self): try: - if await self.client.has_collection(self.collection_name): + collections = await self.client.list_collections() + if self.collection_name in collections: await self.client.drop_collection(collection_name=self.collection_name) except Exception as e: logger.warning(f"Failed to check or delete collection {self.collection_name}: {e}") @@ -78,7 +79,8 @@ class MilvusIndex(EmbeddingIndex): ) try: - collection_exists = await self.client.has_collection(self.collection_name) + collections = await self.client.list_collections() + collection_exists = self.collection_name in collections except Exception as e: logger.error(f"Failed to check collection existence: {self.collection_name} ({e})") # If it's an event loop issue, try to recreate the client @@ -87,7 +89,8 @@ class MilvusIndex(EmbeddingIndex): if hasattr(self, "_parent_adapter"): await self._parent_adapter._recreate_client() - collection_exists = await self.client.has_collection(self.collection_name) + collections = await self.client.list_collections() + collection_exists = self.collection_name in collections else: # Assume collection doesn't exist if we can't check collection_exists = False diff --git a/tests/unit/providers/vector_io/remote/test_milvus.py b/tests/unit/providers/vector_io/remote/test_milvus.py index 25374e617..7020900d7 100644 --- a/tests/unit/providers/vector_io/remote/test_milvus.py +++ b/tests/unit/providers/vector_io/remote/test_milvus.py @@ -40,7 +40,7 @@ async def mock_milvus_client() -> MagicMock: """Create a mock Milvus client with common method behaviors.""" client = MagicMock() - client.has_collection = AsyncMock(return_value=False) # Initially no collection + client.list_collections = AsyncMock(return_value=[]) # Initially no collections client.create_collection = AsyncMock(return_value=None) client.drop_collection = AsyncMock(return_value=None) @@ -103,7 +103,7 @@ async def milvus_index(mock_milvus_client): async def test_add_chunks(milvus_index, sample_chunks, sample_embeddings, mock_milvus_client): # Setup: collection doesn't exist initially, then exists after creation - mock_milvus_client.has_collection.side_effect = [False, True] + mock_milvus_client.list_collections.side_effect = [[], ["test_collection"]] await milvus_index.add_chunks(sample_chunks, sample_embeddings) @@ -120,7 +120,7 @@ async def test_query_chunks_vector( milvus_index, sample_chunks, sample_embeddings, embedding_dimension, mock_milvus_client ): # Setup: Add chunks first - mock_milvus_client.has_collection.return_value = True + mock_milvus_client.list_collections.return_value = ["test_collection"] await milvus_index.add_chunks(sample_chunks, sample_embeddings) # Test vector search @@ -133,7 +133,7 @@ async def test_query_chunks_vector( async def test_query_chunks_keyword_search(milvus_index, sample_chunks, sample_embeddings, mock_milvus_client): - mock_milvus_client.has_collection.return_value = True + mock_milvus_client.list_collections.return_value = ["test_collection"] await milvus_index.add_chunks(sample_chunks, sample_embeddings) # Test keyword search @@ -146,7 +146,7 @@ async def test_query_chunks_keyword_search(milvus_index, sample_chunks, sample_e async def test_bm25_fallback_to_simple_search(milvus_index, sample_chunks, sample_embeddings, mock_milvus_client): """Test that when BM25 search fails, the system falls back to simple text search.""" - mock_milvus_client.has_collection.return_value = True + mock_milvus_client.list_collections.return_value = ["test_collection"] await milvus_index.add_chunks(sample_chunks, sample_embeddings) # Force BM25 search to fail @@ -188,7 +188,7 @@ async def test_bm25_fallback_to_simple_search(milvus_index, sample_chunks, sampl async def test_delete_collection(milvus_index, mock_milvus_client): # Test collection deletion - mock_milvus_client.has_collection.return_value = True + mock_milvus_client.list_collections.return_value = ["test_collection"] await milvus_index.delete() @@ -199,7 +199,7 @@ async def test_query_hybrid_search_rrf( milvus_index, sample_chunks, sample_embeddings, embedding_dimension, mock_milvus_client ): """Test hybrid search with RRF reranker.""" - mock_milvus_client.has_collection.return_value = True + mock_milvus_client.list_collections.return_value = ["test_collection"] await milvus_index.add_chunks(sample_chunks, sample_embeddings) # Mock hybrid search results @@ -251,7 +251,7 @@ async def test_query_hybrid_search_weighted( milvus_index, sample_chunks, sample_embeddings, embedding_dimension, mock_milvus_client ): """Test hybrid search with weighted reranker.""" - mock_milvus_client.has_collection.return_value = True + mock_milvus_client.list_collections.return_value = ["test_collection"] await milvus_index.add_chunks(sample_chunks, sample_embeddings) # Mock hybrid search results @@ -297,7 +297,7 @@ async def test_query_hybrid_search_default_rrf( milvus_index, sample_chunks, sample_embeddings, embedding_dimension, mock_milvus_client ): """Test hybrid search with default RRF reranker (no reranker_type specified).""" - mock_milvus_client.has_collection.return_value = True + mock_milvus_client.list_collections.return_value = ["test_collection"] await milvus_index.add_chunks(sample_chunks, sample_embeddings) # Mock hybrid search results