mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-03 19:57:35 +00:00
refactor(client): replace all AsyncMilvusClient usage of has_collection() with list_collections()
Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>
This commit is contained in:
parent
5482396459
commit
295d8b99c3
2 changed files with 15 additions and 12 deletions
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue