Feature/streaming triples (#676)

* Steaming triples

* Also GraphRAG service uses this

* Updated tests
This commit is contained in:
cybermaggedon 2026-03-09 15:46:33 +00:00 committed by GitHub
parent 3c3e11bef5
commit d2d71f859d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 542 additions and 116 deletions

View file

@ -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,
@ -790,6 +869,74 @@ 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)
# 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,
query: str,