apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-05-11 00:02:37 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Python API implements streaming interfaces (#577)

Browse source 
* Tech spec

* Python CLI utilities updated to use the API including streaming features

* Added type safety to Python API

* Completed missing auth token support in CLI
This commit is contained in:
cybermaggedon 2025-12-04 17:38:57 +00:00 • committed by GitHub
parent b957004db9
commit 01aeede78b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
53 changed files with 4489 additions and 715 deletions
Download patch file Download diff file Expand all files Collapse all files

172
trustgraph-cli/trustgraph/cli/invoke_graph_rag.py
View file

@ -4,13 +4,10 @@ Uses the GraphRAG service to answer a question
import argparse
import os
import asyncio
import json
import uuid
from websockets.asyncio.client import connect
from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
default_user = 'trustgraph'
default_collection = 'default'
default_entity_limit = 50
@ -18,89 +15,51 @@ default_triple_limit = 30
default_max_subgraph_size = 150
default_max_path_length = 2
async def question_streaming(
def question(
url, flow_id, question, user, collection, entity_limit, triple_limit,
max_subgraph_size, max_path_length
max_subgraph_size, max_path_length, streaming=True, token=None
):
"""Streaming version using websockets"""
# Convert http:// to ws://
if url.startswith('http://'):
url = 'ws://' + url[7:]
elif url.startswith('https://'):
url = 'wss://' + url[8:]
# Create API client
api = Api(url=url, token=token)
if not url.endswith("/"):
url += "/"
if streaming:
# Use socket client for streaming
socket = api.socket()
flow = socket.flow(flow_id)
url = url + "api/v1/socket"
try:
response = flow.graph_rag(
question=question,
user=user,
collection=collection,
entity_limit=entity_limit,
triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length,
streaming=True
)
mid = str(uuid.uuid4())
# Stream output
for chunk in response:
print(chunk.content, end="", flush=True)
print() # Final newline
async with connect(url) as ws:
req = {
"id": mid,
"service": "graph-rag",
"flow": flow_id,
"request": {
"query": question,
"user": user,
"collection": collection,
"entity-limit": entity_limit,
"triple-limit": triple_limit,
"max-subgraph-size": max_subgraph_size,
"max-path-length": max_path_length,
"streaming": True
}
}
req = json.dumps(req)
await ws.send(req)
while True:
msg = await ws.recv()
obj = json.loads(msg)
if "error" in obj:
raise RuntimeError(obj["error"])
if obj["id"] != mid:
print("Ignore message")
continue
response = obj["response"]
# Handle streaming format (chunk)
if "chunk" in response:
chunk = response["chunk"]
print(chunk, end="", flush=True)
elif "response" in response:
# Final response with complete text
# Already printed via chunks, just add newline
pass
if obj["complete"]:
print() # Final newline
break
await ws.close()
def question_non_streaming(
url, flow_id, question, user, collection, entity_limit, triple_limit,
max_subgraph_size, max_path_length
):
"""Non-streaming version using HTTP API"""
api = Api(url).flow().id(flow_id)
resp = api.graph_rag(
question=question, user=user, collection=collection,
entity_limit=entity_limit, triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length
)
print(resp)
finally:
socket.close()
else:
# Use REST API for non-streaming
flow = api.flow().id(flow_id)
resp = flow.graph_rag(
question=question,
user=user,
collection=collection,
entity_limit=entity_limit,
triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length
)
print(resp)
def main():
@ -115,6 +74,12 @@ def main():
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument(
'-f', '--flow-id',
default="default",
@ -141,24 +106,28 @@ def main():
parser.add_argument(
'-e', '--entity-limit',
type=int,
default=default_entity_limit,
help=f'Entity limit (default: {default_entity_limit})'
)
parser.add_argument(
'-t', '--triple-limit',
'--triple-limit',
type=int,
default=default_triple_limit,
help=f'Triple limit (default: {default_triple_limit})'
)
parser.add_argument(
'-s', '--max-subgraph-size',
type=int,
default=default_max_subgraph_size,
help=f'Max subgraph size (default: {default_max_subgraph_size})'
)
parser.add_argument(
'-p', '--max-path-length',
type=int,
default=default_max_path_length,
help=f'Max path length (default: {default_max_path_length})'
)
@ -173,36 +142,23 @@ def main():
try:
if not args.no_streaming:
asyncio.run(
question_streaming(
url=args.url,
flow_id=args.flow_id,
question=args.question,
user=args.user,
collection=args.collection,
entity_limit=args.entity_limit,
triple_limit=args.triple_limit,
max_subgraph_size=args.max_subgraph_size,
max_path_length=args.max_path_length,
)
)
else:
question_non_streaming(
url=args.url,
flow_id=args.flow_id,
question=args.question,
user=args.user,
collection=args.collection,
entity_limit=args.entity_limit,
triple_limit=args.triple_limit,
max_subgraph_size=args.max_subgraph_size,
max_path_length=args.max_path_length,
)
question(
url=args.url,
flow_id=args.flow_id,
question=args.question,
user=args.user,
collection=args.collection,
entity_limit=args.entity_limit,
triple_limit=args.triple_limit,
max_subgraph_size=args.max_subgraph_size,
max_path_length=args.max_path_length,
streaming=not args.no_streaming,
token=args.token,
)
except Exception as e:
print("Exception:", e, flush=True)
if __name__ == "__main__":
main()
main()