apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 04:31:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Support for streaming Graph/Doc RAG

Browse source
This commit is contained in:
Cyber MacGeddon 2025-11-26 17:10:52 +00:00
parent 44a170ea28
commit 985c725c50
8 changed files with 418 additions and 90 deletions
Download patch file Download diff file Expand all files Collapse all files

88
trustgraph-base/trustgraph/messaging/translators/retrieval.py
View file

@ -5,43 +5,65 @@ from .base import MessageTranslator
class DocumentRagRequestTranslator(MessageTranslator):
"""Translator for DocumentRagQuery schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> DocumentRagQuery:
return DocumentRagQuery(
query=data["query"],
user=data.get("user", "trustgraph"),
collection=data.get("collection", "default"),
doc_limit=int(data.get("doc-limit", 20))
doc_limit=int(data.get("doc-limit", 20)),
streaming=data.get("streaming", False)
)
def from_pulsar(self, obj: DocumentRagQuery) -> Dict[str, Any]:
return {
"query": obj.query,
"user": obj.user,
"collection": obj.collection,
"doc-limit": obj.doc_limit
"doc-limit": obj.doc_limit,
"streaming": getattr(obj, "streaming", False)
}
class DocumentRagResponseTranslator(MessageTranslator):
"""Translator for DocumentRagResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> DocumentRagResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: DocumentRagResponse) -> Dict[str, Any]:
return {
"response": obj.response
}
result = {}
# Check if this is a streaming response (has chunk)
if hasattr(obj, 'chunk') and obj.chunk:
result["chunk"] = obj.chunk
result["end_of_stream"] = getattr(obj, "end_of_stream", False)
else:
# Non-streaming response
if obj.response:
result["response"] = obj.response
# Always include error if present
if hasattr(obj, 'error') and obj.error and obj.error.message:
result["error"] = {"message": obj.error.message, "type": obj.error.type}
return result
def from_response_with_completion(self, obj: DocumentRagResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True
# For streaming responses, check end_of_stream
if hasattr(obj, 'chunk') and obj.chunk:
is_final = getattr(obj, 'end_of_stream', False)
else:
# For non-streaming responses, it's always final
is_final = True
return self.from_pulsar(obj), is_final
class GraphRagRequestTranslator(MessageTranslator):
"""Translator for GraphRagQuery schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> GraphRagQuery:
return GraphRagQuery(
query=data["query"],
@ -50,9 +72,10 @@ class GraphRagRequestTranslator(MessageTranslator):
entity_limit=int(data.get("entity-limit", 50)),
triple_limit=int(data.get("triple-limit", 30)),
max_subgraph_size=int(data.get("max-subgraph-size", 1000)),
max_path_length=int(data.get("max-path-length", 2))
max_path_length=int(data.get("max-path-length", 2)),
streaming=data.get("streaming", False)
)
def from_pulsar(self, obj: GraphRagQuery) -> Dict[str, Any]:
return {
"query": obj.query,
@ -61,21 +84,42 @@ class GraphRagRequestTranslator(MessageTranslator):
"entity-limit": obj.entity_limit,
"triple-limit": obj.triple_limit,
"max-subgraph-size": obj.max_subgraph_size,
"max-path-length": obj.max_path_length
"max-path-length": obj.max_path_length,
"streaming": getattr(obj, "streaming", False)
}
class GraphRagResponseTranslator(MessageTranslator):
"""Translator for GraphRagResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> GraphRagResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: GraphRagResponse) -> Dict[str, Any]:
return {
"response": obj.response
}
result = {}
# Check if this is a streaming response (has chunk)
if hasattr(obj, 'chunk') and obj.chunk:
result["chunk"] = obj.chunk
result["end_of_stream"] = getattr(obj, "end_of_stream", False)
else:
# Non-streaming response
if obj.response:
result["response"] = obj.response
# Always include error if present
if hasattr(obj, 'error') and obj.error and obj.error.message:
result["error"] = {"message": obj.error.message, "type": obj.error.type}
return result
def from_response_with_completion(self, obj: GraphRagResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True
# For streaming responses, check end_of_stream
if hasattr(obj, 'chunk') and obj.chunk:
is_final = getattr(obj, 'end_of_stream', False)
else:
# For non-streaming responses, it's always final
is_final = True
return self.from_pulsar(obj), is_final

6
trustgraph-base/trustgraph/schema/services/retrieval.py
View file

@ -15,10 +15,13 @@ class GraphRagQuery(Record):
triple_limit = Integer()
max_subgraph_size = Integer()
max_path_length = Integer()
streaming = Boolean()
class GraphRagResponse(Record):
error = Error()
response = String()
chunk = String()
end_of_stream = Boolean()
############################################################################
@ -29,8 +32,11 @@ class DocumentRagQuery(Record):
user = String()
collection = String()
doc_limit = Integer()
streaming = Boolean()
class DocumentRagResponse(Record):
error = Error()
response = String()
chunk = String()
end_of_stream = Boolean()

102
trustgraph-cli/trustgraph/cli/invoke_document_rag.py
View file

@ -4,6 +4,10 @@ Uses the DocumentRAG 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/')
@ -11,7 +15,69 @@ default_user = 'trustgraph'
default_collection = 'default'
default_doc_limit = 10
def question(url, flow_id, question, user, collection, doc_limit):
async def question_streaming(url, flow_id, question, user, collection, doc_limit):
"""Streaming version using websockets"""
# Convert http:// to ws://
if url.startswith('http://'):
url = 'ws://' + url[7:]
elif url.startswith('https://'):
url = 'wss://' + url[8:]
if not url.endswith("/"):
url += "/"
url = url + "api/v1/socket"
mid = str(uuid.uuid4())
async with connect(url) as ws:
req = {
"id": mid,
"service": "document-rag",
"flow": flow_id,
"request": {
"query": question,
"user": user,
"collection": collection,
"doc-limit": doc_limit,
"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, doc_limit):
"""Non-streaming version using HTTP API"""
api = Api(url).flow().id(flow_id)
@ -65,18 +131,36 @@ def main():
help=f'Document limit (default: {default_doc_limit})'
)
parser.add_argument(
'--streaming',
action='store_true',
help='Enable streaming mode (token-by-token output)'
)
args = parser.parse_args()
try:
question(
url=args.url,
flow_id = args.flow_id,
question=args.question,
user=args.user,
collection=args.collection,
doc_limit=args.doc_limit,
)
if args.streaming:
asyncio.run(
question_streaming(
url=args.url,
flow_id=args.flow_id,
question=args.question,
user=args.user,
collection=args.collection,
doc_limit=args.doc_limit,
)
)
else:
question_non_streaming(
url=args.url,
flow_id=args.flow_id,
question=args.question,
user=args.user,
collection=args.collection,
doc_limit=args.doc_limit,
)
except Exception as e:

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

@ -4,6 +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/')
@ -14,10 +18,78 @@ default_triple_limit = 30
default_max_subgraph_size = 150
default_max_path_length = 2
def question(
async def question_streaming(
url, flow_id, question, user, collection, entity_limit, triple_limit,
max_subgraph_size, max_path_length
):
"""Streaming version using websockets"""
# Convert http:// to ws://
if url.startswith('http://'):
url = 'ws://' + url[7:]
elif url.startswith('https://'):
url = 'wss://' + url[8:]
if not url.endswith("/"):
url += "/"
url = url + "api/v1/socket"
mid = str(uuid.uuid4())
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)
@ -91,21 +163,42 @@ def main():
help=f'Max path length (default: {default_max_path_length})'
)
parser.add_argument(
'--streaming',
action='store_true',
help='Enable streaming mode (token-by-token output)'
)
args = parser.parse_args()
try:
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,
)
if args.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,
)
except Exception as e:

18
trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py
View file

@ -68,7 +68,7 @@ class DocumentRag:
async def query(
self, query, user="trustgraph", collection="default",
doc_limit=20,
doc_limit=20, streaming=False, chunk_callback=None,
):
if self.verbose:
@ -86,10 +86,18 @@ class DocumentRag:
logger.debug(f"Documents: {docs}")
logger.debug(f"Query: {query}")
resp = await self.prompt_client.document_prompt(
query = query,
documents = docs
)
if streaming and chunk_callback:
resp = await self.prompt_client.document_prompt(
query=query,
documents=docs,
streaming=True,
chunk_callback=chunk_callback
)
else:
resp = await self.prompt_client.document_prompt(
query=query,
documents=docs
)
if self.verbose:
logger.debug("Query processing complete")

83
trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
View file

@ -92,20 +92,56 @@ class Processor(FlowProcessor):
else:
doc_limit = self.doc_limit
response = await self.rag.query(
v.query,
user=v.user,
collection=v.collection,
doc_limit=doc_limit
)
# Check if streaming is requested
if v.streaming:
# Define async callback for streaming chunks
async def send_chunk(chunk):
await flow("response").send(
DocumentRagResponse(
chunk=chunk,
end_of_stream=False,
response=None,
error=None
),
properties={"id": id}
)
await flow("response").send(
DocumentRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
# Query with streaming enabled
full_response = await self.rag.query(
v.query,
user=v.user,
collection=v.collection,
doc_limit=doc_limit,
streaming=True,
chunk_callback=send_chunk,
)
# Send final message with complete response
await flow("response").send(
DocumentRagResponse(
chunk=None,
end_of_stream=True,
response=full_response,
error=None
),
properties={"id": id}
)
else:
# Non-streaming path (existing behavior)
response = await self.rag.query(
v.query,
user=v.user,
collection=v.collection,
doc_limit=doc_limit
)
await flow("response").send(
DocumentRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
logger.info("Request processing complete")
@ -115,14 +151,21 @@ class Processor(FlowProcessor):
logger.debug("Sending error response...")
await flow("response").send(
DocumentRagResponse(
response = None,
error = Error(
type = "document-rag-error",
message = str(e),
),
# Send error response with end_of_stream flag if streaming was requested
error_response = DocumentRagResponse(
response = None,
error = Error(
type = "document-rag-error",
message = str(e),
),
)
# If streaming was requested, indicate stream end
if v.streaming:
error_response.end_of_stream = True
await flow("response").send(
error_response,
properties = {"id": id}
)

11
trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
View file

@ -316,7 +316,7 @@ class GraphRag:
async def query(
self, query, user = "trustgraph", collection = "default",
entity_limit = 50, triple_limit = 30, max_subgraph_size = 1000,
max_path_length = 2,
max_path_length = 2, streaming = False, chunk_callback = None,
):
if self.verbose:
@ -337,7 +337,14 @@ class GraphRag:
logger.debug(f"Knowledge graph: {kg}")
logger.debug(f"Query: {query}")
resp = await self.prompt_client.kg_prompt(query, kg)
if streaming and chunk_callback:
resp = await self.prompt_client.kg_prompt(
query, kg,
streaming=True,
chunk_callback=chunk_callback
)
else:
resp = await self.prompt_client.kg_prompt(query, kg)
if self.verbose:
logger.debug("Query processing complete")

83
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
View file

@ -135,20 +135,56 @@ class Processor(FlowProcessor):
else:
max_path_length = self.default_max_path_length
response = await rag.query(
query = v.query, user = v.user, collection = v.collection,
entity_limit = entity_limit, triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
)
# Check if streaming is requested
if v.streaming:
# Define async callback for streaming chunks
async def send_chunk(chunk):
await flow("response").send(
GraphRagResponse(
chunk=chunk,
end_of_stream=False,
response=None,
error=None
),
properties={"id": id}
)
await flow("response").send(
GraphRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
# Query with streaming enabled
full_response = await rag.query(
query = v.query, user = v.user, collection = v.collection,
entity_limit = entity_limit, triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
streaming = True,
chunk_callback = send_chunk,
)
# Send final message with complete response
await flow("response").send(
GraphRagResponse(
chunk=None,
end_of_stream=True,
response=full_response,
error=None
),
properties={"id": id}
)
else:
# Non-streaming path (existing behavior)
response = await rag.query(
query = v.query, user = v.user, collection = v.collection,
entity_limit = entity_limit, triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
)
await flow("response").send(
GraphRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
logger.info("Request processing complete")
@ -158,14 +194,21 @@ class Processor(FlowProcessor):
logger.debug("Sending error response...")
await flow("response").send(
GraphRagResponse(
response = None,
error = Error(
type = "graph-rag-error",
message = str(e),
),
# Send error response with end_of_stream flag if streaming was requested
error_response = GraphRagResponse(
response = None,
error = Error(
type = "graph-rag-error",
message = str(e),
),
)
# If streaming was requested, indicate stream end
if v.streaming:
error_response.end_of_stream = True
await flow("response").send(
error_response,
properties = {"id": id}
)