From 606f181c5591f6bcc1906c2aef640f16c1b3b8e7 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 25 Nov 2025 18:02:25 +0000 Subject: [PATCH] Some streaming fixes --- trustgraph-base/trustgraph/base/__init__.py | 2 +- .../trustgraph/base/text_completion_client.py | 41 ++++++++-- .../trustgraph/prompt/template/service.py | 80 ++++++++++++++----- 3 files changed, 95 insertions(+), 28 deletions(-) diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 5a97c220..b329f52e 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/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 diff --git a/trustgraph-base/trustgraph/base/text_completion_client.py b/trustgraph-base/trustgraph/base/text_completion_client.py index aba2fada..ae93e22e 100644 --- a/trustgraph-base/trustgraph/base/text_completion_client.py +++ b/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): - resp = await self.request( + async def text_completion(self, system, prompt, streaming=False, timeout=600): + # 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: - raise RuntimeError(resp.error.message) - - return resp.response + return full_response class TextCompletionClientSpec(RequestResponseSpec): def __init__( diff --git a/trustgraph-flow/trustgraph/prompt/template/service.py b/trustgraph-flow/trustgraph/prompt/template/service.py index 1b6822cc..8ae64391 100755 --- a/trustgraph-flow/trustgraph/prompt/template/service.py +++ b/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):