mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-08 21:02:12 +02:00
Also GraphRAG service uses this
This commit is contained in:
parent
027ef9b85d
commit
9a570036b7
2 changed files with 67 additions and 9 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue