Implement GraphRAG optimisation and update tests

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,7 +232,10 @@ 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(
@@ -243,13 +250,18 @@ class TestQuery:
 
         # 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
 
@@ -279,16 +291,20 @@ class TestQuery:
             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,7 +458,7 @@ 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(
@@ -455,11 +472,11 @@ class TestQuery:
         # Mock get_entities to return test entities
         query.get_entities = AsyncMock(return_value=["entity1", "entity2"])
 
-        # Mock follow_edges to add triples to subgraph
+        # Mock follow_edges_batch to return test triples
-        async def mock_follow_edges(ent, subgraph, path_length):
+        query.follow_edges_batch = AsyncMock(return_value={
-            subgraph.add((ent, "predicate", "object"))
+            ("entity1", "predicate1", "object1"),
-        
+            ("entity2", "predicate2", "object2")
-        query.follow_edges = AsyncMock(side_effect=mock_follow_edges)
+        })
 
         # Call get_subgraph
         result = await query.get_subgraph("test query")
@@ -467,14 +484,14 @@ class TestQuery:
         # Verify get_entities was called
         query.get_entities.assert_called_once_with("test query")
 
-        # Verify follow_edges was called for each entity
+        # Verify follow_edges_batch was called with entities and max_path_length
-        assert query.follow_edges.call_count == 2
+        query.follow_edges_batch.assert_called_once_with(["entity1", "entity2"], 1)
-        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 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):

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:
+        # CRITICAL SECURITY: Cache key MUST include user and collection
-            return self.rag.label_cache[e]
+        # 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)
+        label = str(res[0].o)
-        return self.rag.label_cache[e]
+        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):
-
+        """Legacy method - replaced by follow_edges_batch"""
-        # Not needed?
+        # 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(
+        # For backward compatibility, convert to new approach
-            s=ent, p=None, o=None,
+        batch_result = await self.follow_edges_batch([ent], path_length)
-            limit=self.triple_limit,
+        subgraph.update(batch_result)
-            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
-                )
 
     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:
+        return list(subgraph)
-            await self.follow_edges(ent, subgraph, self.max_path_length)
 
-        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 s, p, o in filtered_subgraph:
-        for edge in subgraph:
+            labeled_triple = (
-
+                label_map.get(s, s),
-            if edge[1] == LABEL:
+                label_map.get(p, p),
-                continue
+                label_map.get(o, o)
-
+            )
-            s = await self.maybe_label(edge[0])
+            sg2.append(labeled_triple)
-            p = await self.maybe_label(edge[1])
-            o = await self.maybe_label(edge[2])
-
-            sg2.append((s, p, o))
 
         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")

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(
+            # CRITICAL SECURITY: Create new GraphRag instance per request
-                embeddings_client = flow("embeddings-request"),
+            # This ensures proper isolation between users and collections
-                graph_embeddings_client = flow("graph-embeddings-request"),
+            # Flow clients are request-scoped and must not be shared
-                triples_client = flow("triples-request"),
+            rag = GraphRag(
-                prompt_client = flow("prompt-request"),
+                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,