From da8d3f23bb358f36e973c74a00989b5847040e68 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 23 Sep 2025 19:04:01 +0100 Subject: [PATCH] Implement GraphRAG optimisation and update tests --- tests/unit/test_retrieval/test_graph_rag.py | 87 +++--- .../retrieval/graph_rag/graph_rag.py | 250 +++++++++++++----- .../trustgraph/retrieval/graph_rag/rag.py | 19 +- 3 files changed, 252 insertions(+), 104 deletions(-) diff --git a/tests/unit/test_retrieval/test_graph_rag.py b/tests/unit/test_retrieval/test_graph_rag.py index 788f71a2..5f54e28a 100644 --- a/tests/unit/test_retrieval/test_graph_rag.py +++ b/tests/unit/test_retrieval/test_graph_rag.py @@ -34,7 +34,9 @@ class TestGraphRag: assert graph_rag.graph_embeddings_client == mock_graph_embeddings_client assert graph_rag.triples_client == mock_triples_client assert graph_rag.verbose is False # Default value - assert graph_rag.label_cache == {} # Empty cache initially + # Verify label_cache is an LRUCacheWithTTL instance + from trustgraph.retrieval.graph_rag.graph_rag import LRUCacheWithTTL + assert isinstance(graph_rag.label_cache, LRUCacheWithTTL) def test_graph_rag_initialization_with_verbose(self): """Test GraphRag initialization with verbose enabled""" @@ -59,7 +61,9 @@ class TestGraphRag: assert graph_rag.graph_embeddings_client == mock_graph_embeddings_client assert graph_rag.triples_client == mock_triples_client assert graph_rag.verbose is True - assert graph_rag.label_cache == {} # Empty cache initially + # Verify label_cache is an LRUCacheWithTTL instance + from trustgraph.retrieval.graph_rag.graph_rag import LRUCacheWithTTL + assert isinstance(graph_rag.label_cache, LRUCacheWithTTL) class TestQuery: @@ -228,8 +232,11 @@ class TestQuery: """Test Query.maybe_label method with cached label""" # Create mock GraphRag with label cache mock_rag = MagicMock() - mock_rag.label_cache = {"entity1": "Entity One Label"} - + # Create mock LRUCacheWithTTL + mock_cache = MagicMock() + mock_cache.get.return_value = "Entity One Label" + mock_rag.label_cache = mock_cache + # Initialize Query query = Query( rag=mock_rag, @@ -237,27 +244,32 @@ class TestQuery: collection="test_collection", verbose=False ) - + # Call maybe_label with cached entity result = await query.maybe_label("entity1") - + # Verify cached label is returned assert result == "Entity One Label" + # Verify cache was checked with proper key format (user:collection:entity) + mock_cache.get.assert_called_once_with("test_user:test_collection:entity1") @pytest.mark.asyncio async def test_maybe_label_with_label_lookup(self): """Test Query.maybe_label method with database label lookup""" # Create mock GraphRag with triples client mock_rag = MagicMock() - mock_rag.label_cache = {} # Empty cache + # Create mock LRUCacheWithTTL that returns None (cache miss) + mock_cache = MagicMock() + mock_cache.get.return_value = None + mock_rag.label_cache = mock_cache mock_triples_client = AsyncMock() mock_rag.triples_client = mock_triples_client - + # Mock triple result with label mock_triple = MagicMock() mock_triple.o = "Human Readable Label" mock_triples_client.query.return_value = [mock_triple] - + # Initialize Query query = Query( rag=mock_rag, @@ -265,10 +277,10 @@ class TestQuery: collection="test_collection", verbose=False ) - + # Call maybe_label result = await query.maybe_label("http://example.com/entity") - + # Verify triples client was called correctly mock_triples_client.query.assert_called_once_with( s="http://example.com/entity", @@ -278,17 +290,21 @@ class TestQuery: user="test_user", collection="test_collection" ) - - # Verify result and cache update + + # Verify result and cache update with proper key assert result == "Human Readable Label" - assert mock_rag.label_cache["http://example.com/entity"] == "Human Readable Label" + cache_key = "test_user:test_collection:http://example.com/entity" + mock_cache.put.assert_called_once_with(cache_key, "Human Readable Label") @pytest.mark.asyncio async def test_maybe_label_with_no_label_found(self): """Test Query.maybe_label method when no label is found""" # Create mock GraphRag with triples client mock_rag = MagicMock() - mock_rag.label_cache = {} # Empty cache + # Create mock LRUCacheWithTTL that returns None (cache miss) + mock_cache = MagicMock() + mock_cache.get.return_value = None + mock_rag.label_cache = mock_cache mock_triples_client = AsyncMock() mock_rag.triples_client = mock_triples_client @@ -318,7 +334,8 @@ class TestQuery: # Verify result is entity itself and cache is updated assert result == "unlabeled_entity" - assert mock_rag.label_cache["unlabeled_entity"] == "unlabeled_entity" + cache_key = "test_user:test_collection:unlabeled_entity" + mock_cache.put.assert_called_once_with(cache_key, "unlabeled_entity") @pytest.mark.asyncio async def test_follow_edges_basic_functionality(self): @@ -441,40 +458,40 @@ class TestQuery: @pytest.mark.asyncio async def test_get_subgraph_method(self): """Test Query.get_subgraph method orchestrates entity and edge discovery""" - # Create mock Query that patches get_entities and follow_edges + # Create mock Query that patches get_entities and follow_edges_batch mock_rag = MagicMock() - + query = Query( rag=mock_rag, - user="test_user", + user="test_user", collection="test_collection", verbose=False, max_path_length=1 ) - + # Mock get_entities to return test entities query.get_entities = AsyncMock(return_value=["entity1", "entity2"]) - - # Mock follow_edges to add triples to subgraph - async def mock_follow_edges(ent, subgraph, path_length): - subgraph.add((ent, "predicate", "object")) - - query.follow_edges = AsyncMock(side_effect=mock_follow_edges) - + + # Mock follow_edges_batch to return test triples + query.follow_edges_batch = AsyncMock(return_value={ + ("entity1", "predicate1", "object1"), + ("entity2", "predicate2", "object2") + }) + # Call get_subgraph result = await query.get_subgraph("test query") - + # Verify get_entities was called query.get_entities.assert_called_once_with("test query") - - # Verify follow_edges was called for each entity - assert query.follow_edges.call_count == 2 - query.follow_edges.assert_any_call("entity1", unittest.mock.ANY, 1) - query.follow_edges.assert_any_call("entity2", unittest.mock.ANY, 1) - - # Verify result is list format + + # Verify follow_edges_batch was called with entities and max_path_length + query.follow_edges_batch.assert_called_once_with(["entity1", "entity2"], 1) + + # Verify result is list format and contains expected triples assert isinstance(result, list) assert len(result) == 2 + assert ("entity1", "predicate1", "object1") in result + assert ("entity2", "predicate2", "object2") in result @pytest.mark.asyncio async def test_get_labelgraph_method(self): diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py index a8b6b244..5f866949 100644 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py @@ -1,12 +1,56 @@ import asyncio import logging +import time +from collections import OrderedDict # Module logger logger = logging.getLogger(__name__) LABEL="http://www.w3.org/2000/01/rdf-schema#label" +class LRUCacheWithTTL: + """LRU cache with TTL for label caching + + CRITICAL SECURITY WARNING: + This cache is shared within a GraphRag instance but GraphRag instances + are created per-request. Cache keys MUST include user:collection prefix + to ensure data isolation between different security contexts. + """ + + def __init__(self, max_size=5000, ttl=300): + self.cache = OrderedDict() + self.access_times = {} + self.max_size = max_size + self.ttl = ttl + + def get(self, key): + if key not in self.cache: + return None + + # Check TTL expiration + if time.time() - self.access_times[key] > self.ttl: + del self.cache[key] + del self.access_times[key] + return None + + # Move to end (most recently used) + self.cache.move_to_end(key) + return self.cache[key] + + def put(self, key, value): + if key in self.cache: + self.cache.move_to_end(key) + else: + if len(self.cache) >= self.max_size: + # Remove least recently used + oldest_key = next(iter(self.cache)) + del self.cache[oldest_key] + del self.access_times[oldest_key] + + self.cache[key] = value + self.access_times[key] = time.time() + class Query: def __init__( @@ -61,8 +105,14 @@ class Query: async def maybe_label(self, e): - if e in self.rag.label_cache: - return self.rag.label_cache[e] + # CRITICAL SECURITY: Cache key MUST include user and collection + # to prevent data leakage between different contexts + cache_key = f"{self.user}:{self.collection}:{e}" + + # Check LRU cache first with isolated key + cached_label = self.rag.label_cache.get(cache_key) + if cached_label is not None: + return cached_label res = await self.rag.triples_client.query( s=e, p=LABEL, o=None, limit=1, @@ -70,60 +120,104 @@ class Query: ) if len(res) == 0: - self.rag.label_cache[e] = e + self.rag.label_cache.put(cache_key, e) return e - self.rag.label_cache[e] = str(res[0].o) - return self.rag.label_cache[e] + label = str(res[0].o) + self.rag.label_cache.put(cache_key, label) + return label + + async def execute_batch_triple_queries(self, entities, limit_per_entity): + """Execute triple queries for multiple entities concurrently""" + tasks = [] + + for entity in entities: + # Create concurrent tasks for all 3 query types per entity + tasks.extend([ + self.rag.triples_client.query( + s=entity, p=None, o=None, + limit=limit_per_entity, + user=self.user, collection=self.collection + ), + self.rag.triples_client.query( + s=None, p=entity, o=None, + limit=limit_per_entity, + user=self.user, collection=self.collection + ), + self.rag.triples_client.query( + s=None, p=None, o=entity, + limit=limit_per_entity, + user=self.user, collection=self.collection + ) + ]) + + # Execute all queries concurrently + results = await asyncio.gather(*tasks, return_exceptions=True) + + # Combine all results + all_triples = [] + for result in results: + if not isinstance(result, Exception): + all_triples.extend(result) + + return all_triples + + async def follow_edges_batch(self, entities, max_depth): + """Optimized iterative graph traversal with batching""" + visited = set() + current_level = set(entities) + subgraph = set() + + for depth in range(max_depth): + if not current_level or len(subgraph) >= self.max_subgraph_size: + break + + # Filter out already visited entities + unvisited_entities = [e for e in current_level if e not in visited] + if not unvisited_entities: + break + + # Batch query all unvisited entities at current level + triples = await self.execute_batch_triple_queries( + unvisited_entities, self.triple_limit + ) + + # Process results and collect next level entities + next_level = set() + for triple in triples: + triple_tuple = (str(triple.s), str(triple.p), str(triple.o)) + subgraph.add(triple_tuple) + + # Collect entities for next level (only from s and o positions) + if depth < max_depth - 1: # Don't collect for final depth + s, p, o = triple_tuple + if s not in visited: + next_level.add(s) + if o not in visited: + next_level.add(o) + + # Stop if subgraph size limit reached + if len(subgraph) >= self.max_subgraph_size: + return subgraph + + # Update for next iteration + visited.update(current_level) + current_level = next_level + + return subgraph async def follow_edges(self, ent, subgraph, path_length): - - # Not needed? + """Legacy method - replaced by follow_edges_batch""" + # Maintain backward compatibility with early termination checks if path_length <= 0: return - # Stop spanning around if the subgraph is already maxed out if len(subgraph) >= self.max_subgraph_size: return - res = await self.rag.triples_client.query( - s=ent, p=None, o=None, - limit=self.triple_limit, - user=self.user, collection=self.collection, - ) - - for triple in res: - subgraph.add( - (str(triple.s), str(triple.p), str(triple.o)) - ) - if path_length > 1: - await self.follow_edges(str(triple.o), subgraph, path_length-1) - - res = await self.rag.triples_client.query( - s=None, p=ent, o=None, - limit=self.triple_limit, - user=self.user, collection=self.collection, - ) - - for triple in res: - subgraph.add( - (str(triple.s), str(triple.p), str(triple.o)) - ) - - res = await self.rag.triples_client.query( - s=None, p=None, o=ent, - limit=self.triple_limit, - user=self.user, collection=self.collection, - ) - - for triple in res: - subgraph.add( - (str(triple.s), str(triple.p), str(triple.o)) - ) - if path_length > 1: - await self.follow_edges( - str(triple.s), subgraph, path_length-1 - ) + # For backward compatibility, convert to new approach + batch_result = await self.follow_edges_batch([ent], path_length) + subgraph.update(batch_result) async def get_subgraph(self, query): @@ -132,31 +226,52 @@ class Query: if self.verbose: logger.debug("Getting subgraph...") - subgraph = set() + # Use optimized batch traversal instead of sequential processing + subgraph = await self.follow_edges_batch(entities, self.max_path_length) - for ent in entities: - await self.follow_edges(ent, subgraph, self.max_path_length) + return list(subgraph) - subgraph = list(subgraph) + async def resolve_labels_batch(self, entities): + """Resolve labels for multiple entities in parallel""" + tasks = [] + for entity in entities: + tasks.append(self.maybe_label(entity)) - return subgraph + return await asyncio.gather(*tasks, return_exceptions=True) async def get_labelgraph(self, query): subgraph = await self.get_subgraph(query) + # Filter out label triples + filtered_subgraph = [edge for edge in subgraph if edge[1] != LABEL] + + # Collect all unique entities that need label resolution + entities_to_resolve = set() + for s, p, o in filtered_subgraph: + entities_to_resolve.update([s, p, o]) + + # Batch resolve labels for all entities in parallel + entity_list = list(entities_to_resolve) + resolved_labels = await self.resolve_labels_batch(entity_list) + + # Create entity-to-label mapping + label_map = {} + for entity, label in zip(entity_list, resolved_labels): + if not isinstance(label, Exception): + label_map[entity] = label + else: + label_map[entity] = entity # Fallback to entity itself + + # Apply labels to subgraph sg2 = [] - - for edge in subgraph: - - if edge[1] == LABEL: - continue - - s = await self.maybe_label(edge[0]) - p = await self.maybe_label(edge[1]) - o = await self.maybe_label(edge[2]) - - sg2.append((s, p, o)) + for s, p, o in filtered_subgraph: + labeled_triple = ( + label_map.get(s, s), + label_map.get(p, p), + label_map.get(o, o) + ) + sg2.append(labeled_triple) sg2 = sg2[0:self.max_subgraph_size] @@ -171,6 +286,13 @@ class Query: return sg2 class GraphRag: + """ + CRITICAL SECURITY: + This class MUST be instantiated per-request to ensure proper isolation + between users and collections. The cache within this instance will only + live for the duration of a single request, preventing cross-contamination + of data between different security contexts. + """ def __init__( self, prompt_client, embeddings_client, graph_embeddings_client, @@ -184,7 +306,9 @@ class GraphRag: self.graph_embeddings_client = graph_embeddings_client self.triples_client = triples_client - self.label_cache = {} + # Replace simple dict with LRU cache with TTL + # CRITICAL: This cache only lives for one request due to per-request instantiation + self.label_cache = LRUCacheWithTTL(max_size=5000, ttl=300) if self.verbose: logger.debug("GraphRag initialized") diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index 4d7b1821..e58f0ac1 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py @@ -45,6 +45,10 @@ class Processor(FlowProcessor): self.default_max_subgraph_size = max_subgraph_size self.default_max_path_length = max_path_length + # CRITICAL SECURITY: NEVER share data between users or collections + # Each user/collection combination MUST have isolated data access + # Caching must NEVER allow information leakage across these boundaries + self.register_specification( ConsumerSpec( name = "request", @@ -93,11 +97,14 @@ class Processor(FlowProcessor): try: - self.rag = GraphRag( - embeddings_client = flow("embeddings-request"), - graph_embeddings_client = flow("graph-embeddings-request"), - triples_client = flow("triples-request"), - prompt_client = flow("prompt-request"), + # CRITICAL SECURITY: Create new GraphRag instance per request + # This ensures proper isolation between users and collections + # Flow clients are request-scoped and must not be shared + rag = GraphRag( + embeddings_client=flow("embeddings-request"), + graph_embeddings_client=flow("graph-embeddings-request"), + triples_client=flow("triples-request"), + prompt_client=flow("prompt-request"), verbose=True, ) @@ -128,7 +135,7 @@ class Processor(FlowProcessor): else: max_path_length = self.default_max_path_length - response = await self.rag.query( + response = await rag.query( query = v.query, user = v.user, collection = v.collection, entity_limit = entity_limit, triple_limit = triple_limit, max_subgraph_size = max_subgraph_size,