Also GraphRAG service uses this

This commit is contained in:
Cyber MacGeddon 2026-03-09 14:56:51 +00:00
parent 027ef9b85d
commit 9a570036b7
2 changed files with 67 additions and 9 deletions

View file

@ -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,

View file

@ -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