Support for streaming Graph/Doc RAG

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 {
+        result = {}
-            "response": obj.response
+
-        }
+        # 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 {
+        result = {}
-            "response": obj.response
+
-        }
+        # 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

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()
 

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(
+        if args.streaming:
-            url=args.url,
+            asyncio.run(
-            flow_id = args.flow_id,
+                question_streaming(
-            question=args.question,
+                    url=args.url,
-            user=args.user,
+                    flow_id=args.flow_id,
-            collection=args.collection,
+                    question=args.question,
-            doc_limit=args.doc_limit,
+                    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:
 

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(
+        if args.streaming:
-            url=args.url,
+            asyncio.run(
-            flow_id = args.flow_id,
+                question_streaming(
-            question=args.question,
+                    url=args.url,
-            user=args.user,
+                    flow_id=args.flow_id,
-            collection=args.collection,
+                    question=args.question,
-            entity_limit=args.entity_limit,
+                    user=args.user,
-            triple_limit=args.triple_limit,
+                    collection=args.collection,
-            max_subgraph_size=args.max_subgraph_size,
+                    entity_limit=args.entity_limit,
-            max_path_length=args.max_path_length,
+                    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:
 

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(
+        if streaming and chunk_callback:
-            query = query,
+            resp = await self.prompt_client.document_prompt(
-            documents = docs
+                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")

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(
+            # Check if streaming is requested
-                v.query, 
+            if v.streaming:
-                user=v.user, 
+                # Define async callback for streaming chunks
-                collection=v.collection, 
+                async def send_chunk(chunk):
-                doc_limit=doc_limit
+                    await flow("response").send(
-            )
+                        DocumentRagResponse(
+                            chunk=chunk,
+                            end_of_stream=False,
+                            response=None,
+                            error=None
+                        ),
+                        properties={"id": id}
+                    )
 
-            await flow("response").send(
+                # Query with streaming enabled
-                DocumentRagResponse(
+                full_response = await self.rag.query(
-                    response = response,
+                    v.query,
-                    error = None
+                    user=v.user,
-                ),
+                    collection=v.collection,
-                properties = {"id": id}
+                    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(
+            # Send error response with end_of_stream flag if streaming was requested
-                DocumentRagResponse(
+            error_response = DocumentRagResponse(
-                    response = None,
+                response = None,
-                    error = Error(
+                error = Error(
-                        type = "document-rag-error",
+                    type = "document-rag-error",
-                        message = str(e),
+                    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}
             )
 

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")

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(
+            # Check if streaming is requested
-                query = v.query, user = v.user, collection = v.collection,
+            if v.streaming:
-                entity_limit = entity_limit, triple_limit = triple_limit,
+                # Define async callback for streaming chunks
-                max_subgraph_size = max_subgraph_size,
+                async def send_chunk(chunk):
-                max_path_length = max_path_length,
+                    await flow("response").send(
-            )
+                        GraphRagResponse(
+                            chunk=chunk,
+                            end_of_stream=False,
+                            response=None,
+                            error=None
+                        ),
+                        properties={"id": id}
+                    )
 
-            await flow("response").send(
+                # Query with streaming enabled
-                GraphRagResponse(
+                full_response = await rag.query(
-                    response = response,
+                    query = v.query, user = v.user, collection = v.collection,
-                    error = None
+                    entity_limit = entity_limit, triple_limit = triple_limit,
-                ),
+                    max_subgraph_size = max_subgraph_size,
-                properties = {"id": id}
+                    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(
+            # Send error response with end_of_stream flag if streaming was requested
-                GraphRagResponse(
+            error_response = GraphRagResponse(
-                    response = None,
+                response = None,
-                    error = Error(
+                error = Error(
-                        type = "graph-rag-error",
+                    type = "graph-rag-error",
-                        message = str(e),
+                    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}
             )