mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-08 21:02:12 +02:00
Steaming triples
This commit is contained in:
parent
3c3e11bef5
commit
027ef9b85d
6 changed files with 226 additions and 54 deletions
|
|
@ -790,6 +790,73 @@ class SocketFlowInstance:
|
|||
|
||||
return self.client._send_request_sync("triples", self.flow_id, request, False)
|
||||
|
||||
def triples_query_stream(
|
||||
self,
|
||||
s: Optional[str] = None,
|
||||
p: Optional[str] = None,
|
||||
o: Optional[str] = None,
|
||||
user: Optional[str] = None,
|
||||
collection: Optional[str] = None,
|
||||
limit: int = 100,
|
||||
batch_size: int = 20,
|
||||
**kwargs: Any
|
||||
) -> Iterator[List[Dict[str, Any]]]:
|
||||
"""
|
||||
Query knowledge graph triples with streaming batches.
|
||||
|
||||
Yields batches of triples as they arrive, reducing time-to-first-result
|
||||
and memory overhead for large result sets.
|
||||
|
||||
Args:
|
||||
s: Subject URI (optional, use None for wildcard)
|
||||
p: Predicate URI (optional, use None for wildcard)
|
||||
o: Object URI or Literal (optional, use None for wildcard)
|
||||
user: User/keyspace identifier (optional)
|
||||
collection: Collection identifier (optional)
|
||||
limit: Maximum results to return (default: 100)
|
||||
batch_size: Triples per batch (default: 20)
|
||||
**kwargs: Additional parameters passed to the service
|
||||
|
||||
Yields:
|
||||
List[Dict]: Batches of triples in wire format
|
||||
|
||||
Example:
|
||||
```python
|
||||
socket = api.socket()
|
||||
flow = socket.flow("default")
|
||||
|
||||
for batch in flow.triples_query_stream(
|
||||
user="trustgraph",
|
||||
collection="default"
|
||||
):
|
||||
for triple in batch:
|
||||
print(triple["s"], triple["p"], triple["o"])
|
||||
```
|
||||
"""
|
||||
request = {
|
||||
"limit": limit,
|
||||
"streaming": True,
|
||||
"batch-size": batch_size,
|
||||
}
|
||||
if s is not None:
|
||||
request["s"] = str(s)
|
||||
if p is not None:
|
||||
request["p"] = str(p)
|
||||
if o is not None:
|
||||
request["o"] = str(o)
|
||||
if user is not None:
|
||||
request["user"] = user
|
||||
if collection is not None:
|
||||
request["collection"] = collection
|
||||
request.update(kwargs)
|
||||
|
||||
for chunk in self.client._send_request_sync("triples", self.flow_id, request, streaming=True):
|
||||
# Each chunk is the raw response containing triples batch
|
||||
if hasattr(chunk, 'response'):
|
||||
yield chunk.response
|
||||
elif isinstance(chunk, dict) and "response" in chunk:
|
||||
yield chunk["response"]
|
||||
|
||||
def rows_query(
|
||||
self,
|
||||
query: str,
|
||||
|
|
|
|||
|
|
@ -52,13 +52,23 @@ class TriplesQueryService(FlowProcessor):
|
|||
|
||||
logger.debug(f"Handling triples query request {id}...")
|
||||
|
||||
triples = await self.query_triples(request)
|
||||
|
||||
logger.debug("Sending triples query response...")
|
||||
r = TriplesQueryResponse(triples=triples, error=None)
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
|
||||
logger.debug("Triples query request completed")
|
||||
if request.streaming:
|
||||
# Streaming mode: send batches
|
||||
async for batch, is_final in self.query_triples_stream(request):
|
||||
r = TriplesQueryResponse(
|
||||
triples=batch,
|
||||
error=None,
|
||||
is_final=is_final,
|
||||
)
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
logger.debug("Triples query streaming completed")
|
||||
else:
|
||||
# Non-streaming mode: single response
|
||||
triples = await self.query_triples(request)
|
||||
logger.debug("Sending triples query response...")
|
||||
r = TriplesQueryResponse(triples=triples, error=None)
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
logger.debug("Triples query request completed")
|
||||
|
||||
except Exception as e:
|
||||
|
||||
|
|
@ -76,6 +86,24 @@ class TriplesQueryService(FlowProcessor):
|
|||
|
||||
await flow("response").send(r, properties={"id": id})
|
||||
|
||||
async def query_triples_stream(self, request):
|
||||
"""
|
||||
Streaming query - yields (batch, is_final) tuples.
|
||||
Default implementation batches results from query_triples.
|
||||
Override for true streaming from backend.
|
||||
"""
|
||||
triples = await self.query_triples(request)
|
||||
batch_size = request.batch_size if request.batch_size > 0 else 20
|
||||
|
||||
for i in range(0, len(triples), batch_size):
|
||||
batch = triples[i:i + batch_size]
|
||||
is_final = (i + batch_size >= len(triples))
|
||||
yield batch, is_final
|
||||
|
||||
# Handle empty result
|
||||
if len(triples) == 0:
|
||||
yield [], True
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
|
|
|
|||
|
|
@ -23,14 +23,18 @@ class TriplesQueryRequestTranslator(MessageTranslator):
|
|||
g=g,
|
||||
limit=int(data.get("limit", 10000)),
|
||||
user=data.get("user", "trustgraph"),
|
||||
collection=data.get("collection", "default")
|
||||
collection=data.get("collection", "default"),
|
||||
streaming=data.get("streaming", False),
|
||||
batch_size=int(data.get("batch-size", 20)),
|
||||
)
|
||||
|
||||
def from_pulsar(self, obj: TriplesQueryRequest) -> Dict[str, Any]:
|
||||
result = {
|
||||
"limit": obj.limit,
|
||||
"user": obj.user,
|
||||
"collection": obj.collection
|
||||
"collection": obj.collection,
|
||||
"streaming": obj.streaming,
|
||||
"batch-size": obj.batch_size,
|
||||
}
|
||||
|
||||
if obj.s:
|
||||
|
|
@ -61,4 +65,4 @@ class TriplesQueryResponseTranslator(MessageTranslator):
|
|||
|
||||
def from_response_with_completion(self, obj: TriplesQueryResponse) -> Tuple[Dict[str, Any], bool]:
|
||||
"""Returns (response_dict, is_final)"""
|
||||
return self.from_pulsar(obj), True
|
||||
return self.from_pulsar(obj), obj.is_final
|
||||
|
|
@ -38,11 +38,14 @@ class TriplesQueryRequest:
|
|||
o: Term | None = None
|
||||
g: str | None = None # Graph IRI. None=default graph, "*"=all graphs
|
||||
limit: int = 0
|
||||
streaming: bool = False # Enable streaming mode (multiple batched responses)
|
||||
batch_size: int = 20 # Triples per batch in streaming mode
|
||||
|
||||
@dataclass
|
||||
class TriplesQueryResponse:
|
||||
error: Error | None = None
|
||||
triples: list[Triple] = field(default_factory=list)
|
||||
is_final: bool = True # False for intermediate batches in streaming mode
|
||||
|
||||
############################################################################
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""
|
||||
Connects to the graph query service and dumps all graph edges in Turtle
|
||||
format with RDF-star support for quoted triples.
|
||||
Uses streaming mode for lower time-to-first-processing.
|
||||
"""
|
||||
|
||||
import rdflib
|
||||
|
|
@ -9,66 +10,82 @@ import sys
|
|||
import argparse
|
||||
import os
|
||||
|
||||
from trustgraph.api import Api, Uri
|
||||
from trustgraph.knowledge import QuotedTriple
|
||||
from trustgraph.api import Api
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_user = 'trustgraph'
|
||||
default_collection = 'default'
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
|
||||
|
||||
def value_to_rdflib(val):
|
||||
"""Convert a TrustGraph value to an rdflib term."""
|
||||
if isinstance(val, Uri):
|
||||
def term_to_rdflib(term):
|
||||
"""Convert a wire-format term to an rdflib term."""
|
||||
if term is None:
|
||||
return None
|
||||
|
||||
t = term.get("t", "")
|
||||
|
||||
if t == "i": # IRI
|
||||
iri = term.get("i", "")
|
||||
# Skip malformed URLs with spaces
|
||||
if " " in val:
|
||||
if " " in iri:
|
||||
return None
|
||||
return rdflib.term.URIRef(val)
|
||||
elif isinstance(val, QuotedTriple):
|
||||
# RDF-star quoted triple
|
||||
s_term = value_to_rdflib(val.s)
|
||||
p_term = value_to_rdflib(val.p)
|
||||
o_term = value_to_rdflib(val.o)
|
||||
return rdflib.term.URIRef(iri)
|
||||
elif t == "l": # Literal
|
||||
value = term.get("v", "")
|
||||
datatype = term.get("d")
|
||||
language = term.get("l")
|
||||
if language:
|
||||
return rdflib.term.Literal(value, lang=language)
|
||||
elif datatype:
|
||||
return rdflib.term.Literal(value, datatype=rdflib.term.URIRef(datatype))
|
||||
else:
|
||||
return rdflib.term.Literal(value)
|
||||
elif t == "r": # Quoted triple (RDF-star)
|
||||
triple = term.get("r", {})
|
||||
s_term = term_to_rdflib(triple.get("s"))
|
||||
p_term = term_to_rdflib(triple.get("p"))
|
||||
o_term = term_to_rdflib(triple.get("o"))
|
||||
if s_term is None or p_term is None or o_term is None:
|
||||
return None
|
||||
# rdflib 6.x+ supports Triple as a term type
|
||||
try:
|
||||
return rdflib.term.Triple((s_term, p_term, o_term))
|
||||
except AttributeError:
|
||||
# Fallback for older rdflib versions - represent as string
|
||||
return rdflib.term.Literal(f"<<{val.s} {val.p} {val.o}>>")
|
||||
# Fallback for older rdflib versions
|
||||
return rdflib.term.Literal(f"<<{s_term} {p_term} {o_term}>>")
|
||||
else:
|
||||
return rdflib.term.Literal(str(val))
|
||||
# Fallback
|
||||
return rdflib.term.Literal(str(term))
|
||||
|
||||
|
||||
def show_graph(url, flow_id, user, collection):
|
||||
def show_graph(url, flow_id, user, collection, limit, batch_size, token=None):
|
||||
|
||||
api = Api(url).flow().id(flow_id)
|
||||
|
||||
rows = api.triples_query(
|
||||
s=None, p=None, o=None,
|
||||
user=user, collection=collection,
|
||||
limit=10_000)
|
||||
socket = Api(url, token=token).socket()
|
||||
flow = socket.flow(flow_id)
|
||||
|
||||
g = rdflib.Graph()
|
||||
|
||||
for row in rows:
|
||||
try:
|
||||
for batch in flow.triples_query_stream(
|
||||
s=None, p=None, o=None,
|
||||
user=user, collection=collection,
|
||||
limit=limit,
|
||||
batch_size=batch_size,
|
||||
):
|
||||
for triple in batch:
|
||||
sv = term_to_rdflib(triple.get("s"))
|
||||
pv = term_to_rdflib(triple.get("p"))
|
||||
ov = term_to_rdflib(triple.get("o"))
|
||||
|
||||
sv = rdflib.term.URIRef(row.s)
|
||||
pv = rdflib.term.URIRef(row.p)
|
||||
ov = value_to_rdflib(row.o)
|
||||
if sv is None or pv is None or ov is None:
|
||||
continue
|
||||
|
||||
if ov is None:
|
||||
continue
|
||||
|
||||
g.add((sv, pv, ov))
|
||||
|
||||
g.serialize(destination="output.ttl", format="turtle")
|
||||
g.add((sv, pv, ov))
|
||||
finally:
|
||||
socket.close()
|
||||
|
||||
buf = io.BytesIO()
|
||||
|
||||
g.serialize(destination=buf, format="turtle")
|
||||
|
||||
sys.stdout.write(buf.getvalue().decode("utf-8"))
|
||||
|
||||
|
||||
|
|
@ -103,6 +120,26 @@ def main():
|
|||
help=f'Collection ID (default: {default_collection})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--token',
|
||||
default=default_token,
|
||||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-l', '--limit',
|
||||
type=int,
|
||||
default=10000,
|
||||
help='Maximum number of triples to return (default: 10000)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-b', '--batch-size',
|
||||
type=int,
|
||||
default=20,
|
||||
help='Triples per streaming batch (default: 20)',
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
|
|
@ -112,6 +149,9 @@ def main():
|
|||
flow_id = args.flow_id,
|
||||
user = args.user,
|
||||
collection = args.collection,
|
||||
limit = args.limit,
|
||||
batch_size = args.batch_size,
|
||||
token = args.token,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""
|
||||
Connects to the graph query service and dumps all graph edges.
|
||||
Uses streaming mode for lower time-to-first-result and reduced memory overhead.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
|
@ -11,17 +12,30 @@ default_user = 'trustgraph'
|
|||
default_collection = 'default'
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
|
||||
def show_graph(url, flow_id, user, collection, token=None):
|
||||
def show_graph(url, flow_id, user, collection, limit, batch_size, token=None):
|
||||
|
||||
api = Api(url, token=token).flow().id(flow_id)
|
||||
socket = Api(url, token=token).socket()
|
||||
flow = socket.flow(flow_id)
|
||||
|
||||
rows = api.triples_query(
|
||||
user=user, collection=collection,
|
||||
s=None, p=None, o=None, limit=10_000,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
print(row.s, row.p, row.o)
|
||||
try:
|
||||
for batch in flow.triples_query_stream(
|
||||
user=user,
|
||||
collection=collection,
|
||||
s=None, p=None, o=None,
|
||||
limit=limit,
|
||||
batch_size=batch_size,
|
||||
):
|
||||
for triple in batch:
|
||||
s = triple.get("s", {})
|
||||
p = triple.get("p", {})
|
||||
o = triple.get("o", {})
|
||||
# Format terms for display
|
||||
s_str = s.get("v", s.get("i", str(s)))
|
||||
p_str = p.get("v", p.get("i", str(p)))
|
||||
o_str = o.get("v", o.get("i", str(o)))
|
||||
print(s_str, p_str, o_str)
|
||||
finally:
|
||||
socket.close()
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -60,6 +74,20 @@ def main():
|
|||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-l', '--limit',
|
||||
type=int,
|
||||
default=10000,
|
||||
help='Maximum number of triples to return (default: 10000)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-b', '--batch-size',
|
||||
type=int,
|
||||
default=20,
|
||||
help='Triples per streaming batch (default: 20)',
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
|
|
@ -69,6 +97,8 @@ def main():
|
|||
flow_id = args.flow_id,
|
||||
user = args.user,
|
||||
collection = args.collection,
|
||||
limit = args.limit,
|
||||
batch_size = args.batch_size,
|
||||
token = args.token,
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue