From 1fb5782c2fbb2fc1afc20aa50c46196d2a3a446f Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 9 Mar 2026 15:24:17 +0000 Subject: [PATCH] Fixing reception --- .../trustgraph/api/socket_client.py | 108 +++++++++++++-- .../query/triples/cassandra/service.py | 126 +++++++++++++++--- 2 files changed, 202 insertions(+), 32 deletions(-) diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index 73f13374..700a4531 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -91,9 +91,18 @@ class SocketClient: service: str, flow: Optional[str], request: Dict[str, Any], - streaming: bool = False - ) -> Union[Dict[str, Any], Iterator[StreamingChunk]]: - """Synchronous wrapper around async WebSocket communication""" + streaming: bool = False, + streaming_raw: bool = False + ) -> Union[Dict[str, Any], Iterator[StreamingChunk], Iterator[Dict[str, Any]]]: + """Synchronous wrapper around async WebSocket communication. + + Args: + service: Service name + flow: Flow ID (optional) + request: Request payload + streaming: Use parsed streaming (for agent/RAG chunk types) + streaming_raw: Use raw streaming (for data batches like triples) + """ # Create event loop if needed try: loop = asyncio.get_event_loop() @@ -105,12 +114,14 @@ class SocketClient: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - if streaming: - # For streaming, we need to return an iterator - # Create a generator that runs async code + if streaming_raw: + # Raw streaming for data batches (triples, rows, etc.) + return self._streaming_generator_raw(service, flow, request, loop) + elif streaming: + # Parsed streaming for agent/RAG chunk types return self._streaming_generator(service, flow, request, loop) else: - # For non-streaming, just run the async code and return result + # Non-streaming single response return loop.run_until_complete(self._send_request_async(service, flow, request)) def _streaming_generator( @@ -120,7 +131,7 @@ class SocketClient: request: Dict[str, Any], loop: asyncio.AbstractEventLoop ) -> Iterator[StreamingChunk]: - """Generator that yields streaming chunks""" + """Generator that yields streaming chunks (for agent/RAG responses)""" async_gen = self._send_request_async_streaming(service, flow, request) try: @@ -137,6 +148,74 @@ class SocketClient: except: pass + def _streaming_generator_raw( + self, + service: str, + flow: Optional[str], + request: Dict[str, Any], + loop: asyncio.AbstractEventLoop + ) -> Iterator[Dict[str, Any]]: + """Generator that yields raw response dicts (for data streaming like triples)""" + async_gen = self._send_request_async_streaming_raw(service, flow, request) + + try: + while True: + try: + data = loop.run_until_complete(async_gen.__anext__()) + yield data + except StopAsyncIteration: + break + finally: + try: + loop.run_until_complete(async_gen.aclose()) + except: + pass + + async def _send_request_async_streaming_raw( + self, + service: str, + flow: Optional[str], + request: Dict[str, Any] + ) -> Iterator[Dict[str, Any]]: + """Async streaming that yields raw response dicts without parsing. + + Used for data streaming (triples, rows, etc.) where responses are + just batches of data, not agent/RAG chunk types. + """ + with self._lock: + self._request_counter += 1 + request_id = f"req-{self._request_counter}" + + ws_url = f"{self.url}/api/v1/socket" + if self.token: + ws_url = f"{ws_url}?token={self.token}" + + message = { + "id": request_id, + "service": service, + "request": request + } + if flow: + message["flow"] = flow + + async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket: + await websocket.send(json.dumps(message)) + + async for raw_message in websocket: + response = json.loads(raw_message) + + if response.get("id") != request_id: + continue + + if "error" in response: + raise_from_error_dict(response["error"]) + + if "response" in response: + yield response["response"] + + if response.get("complete"): + break + async def _send_request_async( self, service: str, @@ -850,12 +929,13 @@ class SocketFlowInstance: 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"] + # Use raw streaming - yields response dicts directly without parsing + for response in self.client._send_request_sync("triples", self.flow_id, request, streaming_raw=True): + # Response is {"response": [...triples...]} from translator + if isinstance(response, dict) and "response" in response: + yield response["response"] + else: + yield response def rows_query( self, diff --git a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py index 4d4290b1..9cea4f48 100755 --- a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py +++ b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py @@ -7,6 +7,7 @@ null. Output is a list of quads. import logging import json +from cassandra.query import SimpleStatement from .... direct.cassandra_kg import ( EntityCentricKnowledgeGraph, GRAPH_WILDCARD, DEFAULT_GRAPH @@ -144,28 +145,30 @@ class Processor(TriplesQueryService): self.cassandra_password = password self.table = None + def ensure_connection(self, user): + """Ensure we have a connection to the correct keyspace.""" + if user != self.table: + KGClass = EntityCentricKnowledgeGraph + + if self.cassandra_username and self.cassandra_password: + self.tg = KGClass( + hosts=self.cassandra_host, + keyspace=user, + username=self.cassandra_username, + password=self.cassandra_password + ) + else: + self.tg = KGClass( + hosts=self.cassandra_host, + keyspace=user, + ) + self.table = user + async def query_triples(self, query): try: - user = query.user - - if user != self.table: - # Use factory function to select implementation - KGClass = EntityCentricKnowledgeGraph - - if self.cassandra_username and self.cassandra_password: - self.tg = KGClass( - hosts=self.cassandra_host, - keyspace=query.user, - username=self.cassandra_username, password=self.cassandra_password - ) - else: - self.tg = KGClass( - hosts=self.cassandra_host, - keyspace=query.user, - ) - self.table = user + self.ensure_connection(query.user) # Extract values from query s_val = get_term_value(query.s) @@ -291,6 +294,93 @@ class Processor(TriplesQueryService): logger.error(f"Exception querying triples: {e}", exc_info=True) raise e + async def query_triples_stream(self, query): + """ + Streaming query - yields (batch, is_final) tuples. + Uses Cassandra's paging to fetch results incrementally. + """ + try: + self.ensure_connection(query.user) + + batch_size = query.batch_size if query.batch_size > 0 else 20 + limit = query.limit if query.limit > 0 else 10000 + + # Extract query pattern + s_val = get_term_value(query.s) + p_val = get_term_value(query.p) + o_val = get_term_value(query.o) + g_val = query.g + + # Helper to extract object metadata from result row + def get_o_metadata(t): + otype = getattr(t, 'otype', None) + dtype = getattr(t, 'dtype', None) + lang = getattr(t, 'lang', None) + return otype, dtype, lang + + # For streaming, we need to execute with fetch_size + # Use the collection table for get_all queries (most common streaming case) + + # Determine which query to use based on pattern + if s_val is None and p_val is None and o_val is None: + # Get all - use collection table with paging + cql = f"SELECT d, s, p, o, otype, dtype, lang FROM {self.tg.collection_table} WHERE collection = %s" + params = [query.collection] + else: + # For specific patterns, fall back to non-streaming + # (these typically return small result sets anyway) + async for batch, is_final in self._fallback_stream(query, batch_size): + yield batch, is_final + return + + # Create statement with fetch_size for true streaming + statement = SimpleStatement(cql, fetch_size=batch_size) + result_set = self.tg.session.execute(statement, params) + + batch = [] + count = 0 + + for row in result_set: + if count >= limit: + break + + g = row.d if hasattr(row, 'd') else DEFAULT_GRAPH + otype, dtype, lang = get_o_metadata(row) + + triple = Triple( + s=create_term(row.s), + p=create_term(row.p), + o=create_term(row.o, otype=otype, dtype=dtype, lang=lang), + g=g if g != DEFAULT_GRAPH else None + ) + batch.append(triple) + count += 1 + + # Yield batch when full (never mark as final mid-stream) + if len(batch) >= batch_size: + yield batch, False + batch = [] + + # Always yield final batch to signal completion + # This handles: remaining rows, empty result set, or exact batch boundary + yield batch, True + + except Exception as e: + logger.error(f"Exception in streaming query: {e}", exc_info=True) + raise e + + async def _fallback_stream(self, query, batch_size): + """Fallback to non-streaming query with post-hoc batching.""" + triples = await self.query_triples(query) + + 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 + + if len(triples) == 0: + yield [], True + @staticmethod def add_args(parser):