From 9a570036b798da39cb79786b537b9e78b9aa82b8 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Mon, 9 Mar 2026 14:56:51 +0000
Subject: [PATCH] Also GraphRAG service uses this

---
 .../trustgraph/base/triples_client.py         | 55 +++++++++++++++++++
 .../retrieval/graph_rag/graph_rag.py          | 21 ++++---
 2 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/trustgraph-base/trustgraph/base/triples_client.py b/trustgraph-base/trustgraph/base/triples_client.py
index 7258d3ca..d9715c09 100644
--- a/trustgraph-base/trustgraph/base/triples_client.py
+++ b/trustgraph-base/trustgraph/base/triples_client.py
@@ -57,6 +57,61 @@ class TriplesClient(RequestResponse):
 
         return triples
 
+    async def query_stream(self, s=None, p=None, o=None, limit=20,
+                           user="trustgraph", collection="default",
+                           batch_size=20, timeout=30,
+                           batch_callback=None):
+        """
+        Streaming triple query - calls callback for each batch as it arrives.
+
+        Args:
+            s, p, o: Triple pattern (None for wildcard)
+            limit: Maximum total triples to return
+            user: User/keyspace
+            collection: Collection name
+            batch_size: Triples per batch
+            timeout: Request timeout in seconds
+            batch_callback: Async callback(batch, is_final) called for each batch
+
+        Returns:
+            List[Triple]: All triples (flattened) if no callback provided
+        """
+        all_triples = []
+
+        async def recipient(resp):
+            if resp.error:
+                raise RuntimeError(resp.error.message)
+
+            batch = [
+                Triple(to_value(v.s), to_value(v.p), to_value(v.o))
+                for v in resp.triples
+            ]
+
+            if batch_callback:
+                await batch_callback(batch, resp.is_final)
+            else:
+                all_triples.extend(batch)
+
+            return resp.is_final
+
+        await self.request(
+            TriplesQueryRequest(
+                s=from_value(s),
+                p=from_value(p),
+                o=from_value(o),
+                limit=limit,
+                user=user,
+                collection=collection,
+                streaming=True,
+                batch_size=batch_size,
+            ),
+            timeout=timeout,
+            recipient=recipient,
+        )
+
+        if not batch_callback:
+            return all_triples
+
 class TriplesClientSpec(RequestResponseSpec):
     def __init__(
             self, request_name, response_name,
diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
index 2bf6b2ea..80042951 100644
--- a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
+++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
@@ -129,26 +129,29 @@ class Query:
         return label
 
     async def execute_batch_triple_queries(self, entities, limit_per_entity):
-        """Execute triple queries for multiple entities concurrently"""
+        """Execute triple queries for multiple entities concurrently using streaming"""
         tasks = []
 
         for entity in entities:
-            # Create concurrent tasks for all 3 query types per entity
+            # Create concurrent streaming tasks for all 3 query types per entity
             tasks.extend([
-                self.rag.triples_client.query(
+                self.rag.triples_client.query_stream(
                     s=entity, p=None, o=None,
                     limit=limit_per_entity,
-                    user=self.user, collection=self.collection
+                    user=self.user, collection=self.collection,
+                    batch_size=20,
                 ),
-                self.rag.triples_client.query(
+                self.rag.triples_client.query_stream(
                     s=None, p=entity, o=None,
                     limit=limit_per_entity,
-                    user=self.user, collection=self.collection
+                    user=self.user, collection=self.collection,
+                    batch_size=20,
                 ),
-                self.rag.triples_client.query(
+                self.rag.triples_client.query_stream(
                     s=None, p=None, o=entity,
                     limit=limit_per_entity,
-                    user=self.user, collection=self.collection
+                    user=self.user, collection=self.collection,
+                    batch_size=20,
                 )
             ])
 
@@ -158,7 +161,7 @@ class Query:
         # Combine all results
         all_triples = []
         for result in results:
-            if not isinstance(result, Exception):
+            if not isinstance(result, Exception) and result is not None:
                 all_triples.extend(result)
 
         return all_triples