From 985c725c50f8ae2b5ff0951e55a3e1449906a9aa Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 26 Nov 2025 17:10:52 +0000 Subject: [PATCH] Support for streaming Graph/Doc RAG --- .../messaging/translators/retrieval.py | 88 +++++++++---- .../trustgraph/schema/services/retrieval.py | 6 + .../trustgraph/cli/invoke_document_rag.py | 102 +++++++++++++-- .../trustgraph/cli/invoke_graph_rag.py | 117 ++++++++++++++++-- .../retrieval/document_rag/document_rag.py | 18 ++- .../trustgraph/retrieval/document_rag/rag.py | 83 ++++++++++--- .../retrieval/graph_rag/graph_rag.py | 11 +- .../trustgraph/retrieval/graph_rag/rag.py | 83 ++++++++++--- 8 files changed, 418 insertions(+), 90 deletions(-) diff --git a/trustgraph-base/trustgraph/messaging/translators/retrieval.py b/trustgraph-base/trustgraph/messaging/translators/retrieval.py index 96c25ed8..441a9d18 100644 --- a/trustgraph-base/trustgraph/messaging/translators/retrieval.py +++ b/trustgraph-base/trustgraph/messaging/translators/retrieval.py @@ -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 \ No newline at end of file + # 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 \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/services/retrieval.py b/trustgraph-base/trustgraph/schema/services/retrieval.py index ee96bb1e..3cd7f792 100644 --- a/trustgraph-base/trustgraph/schema/services/retrieval.py +++ b/trustgraph-base/trustgraph/schema/services/retrieval.py @@ -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() diff --git a/trustgraph-cli/trustgraph/cli/invoke_document_rag.py b/trustgraph-cli/trustgraph/cli/invoke_document_rag.py index 8f8c627c..e500da93 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_document_rag.py +++ b/trustgraph-cli/trustgraph/cli/invoke_document_rag.py @@ -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: diff --git a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py b/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py index cf7c64be..accabaa8 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py +++ b/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py @@ -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: diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py index d885757e..9f4ad0ff 100644 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py +++ b/trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py @@ -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") diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py index 2e5149c9..670d71a1 100755 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py @@ -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} ) diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py index 5f866949..7ccba248 100644 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py @@ -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") diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index e58f0ac1..565921a3 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py @@ -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} )