diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index 26c241a7..73f13374 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -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, diff --git a/trustgraph-base/trustgraph/base/triples_query_service.py b/trustgraph-base/trustgraph/base/triples_query_service.py index b156ef55..b8053b01 100644 --- a/trustgraph-base/trustgraph/base/triples_query_service.py +++ b/trustgraph-base/trustgraph/base/triples_query_service.py @@ -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): diff --git a/trustgraph-base/trustgraph/messaging/translators/triples.py b/trustgraph-base/trustgraph/messaging/translators/triples.py index 2b01b1bc..2f29aa56 100644 --- a/trustgraph-base/trustgraph/messaging/translators/triples.py +++ b/trustgraph-base/trustgraph/messaging/translators/triples.py @@ -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 \ No newline at end of file + return self.from_pulsar(obj), obj.is_final \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/services/query.py b/trustgraph-base/trustgraph/schema/services/query.py index 67caa2be..7a65f775 100644 --- a/trustgraph-base/trustgraph/schema/services/query.py +++ b/trustgraph-base/trustgraph/schema/services/query.py @@ -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 ############################################################################ diff --git a/trustgraph-cli/trustgraph/cli/graph_to_turtle.py b/trustgraph-cli/trustgraph/cli/graph_to_turtle.py index f42fe140..840f8574 100644 --- a/trustgraph-cli/trustgraph/cli/graph_to_turtle.py +++ b/trustgraph-cli/trustgraph/cli/graph_to_turtle.py @@ -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: diff --git a/trustgraph-cli/trustgraph/cli/show_graph.py b/trustgraph-cli/trustgraph/cli/show_graph.py index b5b15e3c..105fe604 100644 --- a/trustgraph-cli/trustgraph/cli/show_graph.py +++ b/trustgraph-cli/trustgraph/cli/show_graph.py @@ -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, )