From 4582abd5fc4e1f15436b55eac0043e7b7b71d9e3 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 14 Jul 2025 22:15:19 +0100 Subject: [PATCH] Fixing storage and adding tests --- .../test_query/test_graph_embeddings_milvus_query.py | 10 ++-------- .../query/graph_embeddings/milvus/service.py | 4 ++++ 2 files changed, 6 insertions(+), 8 deletions(-) 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)