Some streaming fixes

trustgraph-base/trustgraph/base/__init__.py

@@ -12,7 +12,7 @@ from . parameter_spec import ParameterSpec
 from . producer_spec import ProducerSpec
 from . subscriber_spec import SubscriberSpec
 from . request_response_spec import RequestResponseSpec
-from . llm_service import LlmService, LlmResult
+from . llm_service import LlmService, LlmResult, LlmChunk
 from . chunking_service import ChunkingService
 from . embeddings_service import EmbeddingsService
 from . embeddings_client import EmbeddingsClientSpec

trustgraph-base/trustgraph/base/text_completion_client.py

@@ -3,18 +3,45 @@ from . request_response_spec import RequestResponse, RequestResponseSpec
 from .. schema import TextCompletionRequest, TextCompletionResponse
 
 class TextCompletionClient(RequestResponse):
-    async def text_completion(self, system, prompt, timeout=600):
+    async def text_completion(self, system, prompt, streaming=False, timeout=600):
-        resp = await self.request(
+        # If not streaming, use original behavior
+        if not streaming:
+            resp = await self.request(
+                TextCompletionRequest(
+                    system = system, prompt = prompt, streaming = False
+                ),
+                timeout=timeout
+            )
+
+            if resp.error:
+                raise RuntimeError(resp.error.message)
+
+            return resp.response
+
+        # For streaming: collect all chunks and return complete response
+        full_response = ""
+
+        async def collect_chunks(resp):
+            nonlocal full_response
+
+            if resp.error:
+                raise RuntimeError(resp.error.message)
+
+            if resp.response:
+                full_response += resp.response
+
+            # Return True when end_of_stream is reached
+            return getattr(resp, 'end_of_stream', False)
+
+        await self.request(
             TextCompletionRequest(
-                system = system, prompt = prompt
+                system = system, prompt = prompt, streaming = True
             ),
+            recipient=collect_chunks,
             timeout=timeout
         )
 
-        if resp.error:
+        return full_response
-            raise RuntimeError(resp.error.message)
-
-        return resp.response
 
 class TextCompletionClientSpec(RequestResponseSpec):
     def __init__(

trustgraph-flow/trustgraph/prompt/template/service.py

@@ -101,6 +101,9 @@ class Processor(FlowProcessor):
 
         kind = v.id
 
+        # Check if streaming is requested
+        streaming = getattr(v, 'streaming', False)
+
         try:
 
             logger.debug(f"Prompt terms: {v.terms}")
@@ -109,16 +112,65 @@ class Processor(FlowProcessor):
                 k: json.loads(v)
                 for k, v in v.terms.items()
             }
-            
-            logger.debug(f"Handling prompt kind {kind}...")
 
+            logger.debug(f"Handling prompt kind {kind}... (streaming={streaming})")
+
+            # If streaming, we need to handle it differently
+            if streaming:
+                # For streaming, we need to intercept LLM responses
+                # and forward them as they arrive
+
+                async def llm_streaming(system, prompt):
+                    logger.debug(f"System prompt: {system}")
+                    logger.debug(f"User prompt: {prompt}")
+
+                    # Use the text completion client with recipient handler
+                    client = flow("text-completion-request")
+
+                    async def forward_chunks(resp):
+                        if resp.error:
+                            raise RuntimeError(resp.error.message)
+
+                        if resp.response:
+                            # Forward each chunk immediately
+                            r = PromptResponse(
+                                text=resp.response,
+                                object=None,
+                                error=None,
+                                end_of_stream=getattr(resp, 'end_of_stream', False),
+                            )
+                            await flow("response").send(r, properties={"id": id})
+
+                        # Return True when end_of_stream
+                        return getattr(resp, 'end_of_stream', False)
+
+                    await client.request(
+                        TextCompletionRequest(
+                            system=system, prompt=prompt, streaming=True
+                        ),
+                        recipient=forward_chunks,
+                        timeout=600
+                    )
+
+                    # Return empty string since we already sent all chunks
+                    return ""
+
+                try:
+                    await self.manager.invoke(kind, input, llm_streaming)
+                except Exception as e:
+                    logger.error(f"Prompt streaming exception: {e}", exc_info=True)
+                    raise e
+
+                return
+
+            # Non-streaming path (original behavior)
             async def llm(system, prompt):
 
                 logger.debug(f"System prompt: {system}")
                 logger.debug(f"User prompt: {prompt}")
 
                 resp = await flow("text-completion-request").text_completion(
-                    system = system, prompt = prompt,
+                    system = system, prompt = prompt, streaming = False,
                 )
 
                 try:
@@ -143,6 +195,7 @@ class Processor(FlowProcessor):
                     text=resp,
                     object=None,
                     error=None,
+                    end_of_stream=True,
                 )
 
                 await flow("response").send(r, properties={"id": id})
@@ -158,6 +211,7 @@ class Processor(FlowProcessor):
                     text=None,
                     object=json.dumps(resp),
                     error=None,
+                    end_of_stream=True,
                 )
 
                 await flow("response").send(r, properties={"id": id})
@@ -175,27 +229,13 @@ class Processor(FlowProcessor):
                     type = "llm-error",
                     message = str(e),
                 ),
-                response=None,
+                text=None,
+                object=None,
+                end_of_stream=True,
             )
 
             await flow("response").send(r, properties={"id": id})
 
-        except Exception as e:
-
-            logger.error(f"Prompt service exception: {e}", exc_info=True)
-
-            logger.debug("Sending error response...")
-
-            r = PromptResponse(
-                error=Error(
-                    type = "llm-error",
-                    message = str(e),
-                ),
-                response=None,
-            )
-
-            await self.send(r, properties={"id": id})
-
     @staticmethod
     def add_args(parser):