diff --git a/tests/unit/test_query/test_graph_embeddings_milvus_query.py b/tests/unit/test_query/test_graph_embeddings_milvus_query.py index 6d856c96..5fbb74d5 100644 --- a/tests/unit/test_query/test_graph_embeddings_milvus_query.py +++ b/tests/unit/test_query/test_graph_embeddings_milvus_query.py @@ -443,16 +443,10 @@ class TestMilvusGraphEmbeddingsQueryProcessor: limit=0 ) - # Mock search results - mock_results = [ - {"entity": {"entity": "http://example.com/entity1"}}, - ] - processor.vecstore.search.return_value = mock_results - result = await processor.query_graph_embeddings(query) - # Verify search was called with 0 limit - processor.vecstore.search.assert_called_once_with([0.1, 0.2, 0.3], limit=0) + # Verify no search was called (optimization for zero limit) + processor.vecstore.search.assert_not_called() # Verify empty results due to zero limit assert len(result) == 0 diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py index 498ab483..7603f4d6 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py @@ -39,6 +39,10 @@ class Processor(GraphEmbeddingsQueryService): entity_set = set() entities = [] + # Handle zero limit case + if msg.limit <= 0: + return [] + for vec in msg.vectors: resp = self.vecstore.search(vec, limit=msg.limit * 2)