From cab6e95d804c49ac27dc5f2f9aac690d44665c66 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 25 Nov 2025 17:15:49 +0000 Subject: [PATCH] Phase 1 & 2 of streaming, covers some VertexAI prototyping --- .../trustgraph/base/llm_service.py | 97 +++++++++++--- .../trustgraph/clients/llm_client.py | 65 +++++++++- .../messaging/translators/prompt.py | 9 +- .../messaging/translators/text_completion.py | 9 +- .../trustgraph/schema/services/llm.py | 4 +- .../trustgraph/schema/services/prompt.py | 8 +- trustgraph-cli/trustgraph/cli/invoke_llm.py | 71 +++++++++-- .../trustgraph/cli/invoke_prompt.py | 85 +++++++++++-- .../model/text_completion/vertexai/llm.py | 119 +++++++++++++++++- 9 files changed, 417 insertions(+), 50 deletions(-) diff --git a/trustgraph-base/trustgraph/base/llm_service.py b/trustgraph-base/trustgraph/base/llm_service.py index 3f5dac43..dc6a8e65 100644 --- a/trustgraph-base/trustgraph/base/llm_service.py +++ b/trustgraph-base/trustgraph/base/llm_service.py @@ -28,6 +28,19 @@ class LlmResult: self.model = model __slots__ = ["text", "in_token", "out_token", "model"] +class LlmChunk: + """Represents a streaming chunk from an LLM""" + def __init__( + self, text = None, in_token = None, out_token = None, + model = None, is_final = False, + ): + self.text = text + self.in_token = in_token + self.out_token = out_token + self.model = model + self.is_final = is_final + __slots__ = ["text", "in_token", "out_token", "model", "is_final"] + class LlmService(FlowProcessor): def __init__(self, **params): @@ -99,16 +112,57 @@ class LlmService(FlowProcessor): id = msg.properties()["id"] - with __class__.text_completion_metric.labels( - id=self.id, - flow=f"{flow.name}-{consumer.name}", - ).time(): + model = flow("model") + temperature = flow("temperature") - model = flow("model") - temperature = flow("temperature") + # Check if streaming is requested and supported + streaming = getattr(request, 'streaming', False) - response = await self.generate_content( - request.system, request.prompt, model, temperature + if streaming and self.supports_streaming(): + + # Streaming mode + with __class__.text_completion_metric.labels( + id=self.id, + flow=f"{flow.name}-{consumer.name}", + ).time(): + + async for chunk in self.generate_content_stream( + request.system, request.prompt, model, temperature + ): + await flow("response").send( + TextCompletionResponse( + error=None, + response=chunk.text, + in_token=chunk.in_token, + out_token=chunk.out_token, + model=chunk.model, + end_of_stream=chunk.is_final + ), + properties={"id": id} + ) + + else: + + # Non-streaming mode (original behavior) + with __class__.text_completion_metric.labels( + id=self.id, + flow=f"{flow.name}-{consumer.name}", + ).time(): + + response = await self.generate_content( + request.system, request.prompt, model, temperature + ) + + await flow("response").send( + TextCompletionResponse( + error=None, + response=response.text, + in_token=response.in_token, + out_token=response.out_token, + model=response.model, + end_of_stream=True + ), + properties={"id": id} ) __class__.text_completion_model_metric.labels( @@ -119,17 +173,6 @@ class LlmService(FlowProcessor): "temperature": str(temperature) if temperature is not None else "", }) - await flow("response").send( - TextCompletionResponse( - error=None, - response=response.text, - in_token=response.in_token, - out_token=response.out_token, - model=response.model - ), - properties={"id": id} - ) - except TooManyRequests as e: raise e @@ -151,10 +194,26 @@ class LlmService(FlowProcessor): in_token=None, out_token=None, model=None, + end_of_stream=True ), properties={"id": id} ) + def supports_streaming(self): + """ + Override in subclass to indicate streaming support. + Returns False by default. + """ + return False + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """ + Override in subclass to implement streaming. + Should yield LlmChunk objects. + The final chunk should have is_final=True. + """ + raise NotImplementedError("Streaming not implemented for this provider") + @staticmethod def add_args(parser): diff --git a/trustgraph-base/trustgraph/clients/llm_client.py b/trustgraph-base/trustgraph/clients/llm_client.py index a8894c8f..3c629e7d 100644 --- a/trustgraph-base/trustgraph/clients/llm_client.py +++ b/trustgraph-base/trustgraph/clients/llm_client.py @@ -5,6 +5,7 @@ from .. schema import TextCompletionRequest, TextCompletionResponse from .. schema import text_completion_request_queue from .. schema import text_completion_response_queue from . base import BaseClient +from .. exceptions import LlmError # Ugly ERROR=_pulsar.LoggerLevel.Error @@ -37,8 +38,68 @@ class LlmClient(BaseClient): output_schema=TextCompletionResponse, ) - def request(self, system, prompt, timeout=300): + def request(self, system, prompt, timeout=300, streaming=False): + """ + Non-streaming request (backward compatible). + Returns complete response string. + """ + if streaming: + raise ValueError("Use request_stream() for streaming requests") return self.call( - system=system, prompt=prompt, timeout=timeout + system=system, prompt=prompt, streaming=False, timeout=timeout ).response + def request_stream(self, system, prompt, timeout=300): + """ + Streaming request generator. + Yields response chunks as they arrive. + Usage: + for chunk in client.request_stream(system, prompt): + print(chunk.response, end='', flush=True) + """ + import time + import uuid + + id = str(uuid.uuid4()) + request = TextCompletionRequest( + system=system, prompt=prompt, streaming=True + ) + + end_time = time.time() + timeout + self.producer.send(request, properties={"id": id}) + + # Collect responses until end_of_stream + while time.time() < end_time: + try: + msg = self.consumer.receive(timeout_millis=2500) + except Exception: + continue + + mid = msg.properties()["id"] + + if mid == id: + value = msg.value() + + # Handle errors + if value.error: + self.consumer.acknowledge(msg) + if value.error.type == "llm-error": + raise LlmError(value.error.message) + else: + raise RuntimeError( + f"{value.error.type}: {value.error.message}" + ) + + self.consumer.acknowledge(msg) + yield value + + # Check if this is the final chunk + if getattr(value, 'end_of_stream', True): + break + else: + # Ignore messages with wrong ID + self.consumer.acknowledge(msg) + + if time.time() >= end_time: + raise TimeoutError("Timed out waiting for response") + diff --git a/trustgraph-base/trustgraph/messaging/translators/prompt.py b/trustgraph-base/trustgraph/messaging/translators/prompt.py index b0e7351f..8916a77c 100644 --- a/trustgraph-base/trustgraph/messaging/translators/prompt.py +++ b/trustgraph-base/trustgraph/messaging/translators/prompt.py @@ -16,10 +16,11 @@ class PromptRequestTranslator(MessageTranslator): k: json.dumps(v) for k, v in data["variables"].items() } - + return PromptRequest( id=data.get("id"), - terms=terms + terms=terms, + streaming=data.get("streaming", False) ) def from_pulsar(self, obj: PromptRequest) -> Dict[str, Any]: @@ -51,4 +52,6 @@ class PromptResponseTranslator(MessageTranslator): def from_response_with_completion(self, obj: PromptResponse) -> Tuple[Dict[str, Any], bool]: """Returns (response_dict, is_final)""" - return self.from_pulsar(obj), True \ No newline at end of file + # Check end_of_stream field to determine if this is the final message + is_final = getattr(obj, 'end_of_stream', True) + return self.from_pulsar(obj), is_final \ No newline at end of file diff --git a/trustgraph-base/trustgraph/messaging/translators/text_completion.py b/trustgraph-base/trustgraph/messaging/translators/text_completion.py index eda3be5d..b4ba4d13 100644 --- a/trustgraph-base/trustgraph/messaging/translators/text_completion.py +++ b/trustgraph-base/trustgraph/messaging/translators/text_completion.py @@ -5,11 +5,12 @@ from .base import MessageTranslator class TextCompletionRequestTranslator(MessageTranslator): """Translator for TextCompletionRequest schema objects""" - + def to_pulsar(self, data: Dict[str, Any]) -> TextCompletionRequest: return TextCompletionRequest( system=data["system"], - prompt=data["prompt"] + prompt=data["prompt"], + streaming=data.get("streaming", False) ) def from_pulsar(self, obj: TextCompletionRequest) -> Dict[str, Any]: @@ -39,4 +40,6 @@ class TextCompletionResponseTranslator(MessageTranslator): def from_response_with_completion(self, obj: TextCompletionResponse) -> Tuple[Dict[str, Any], bool]: """Returns (response_dict, is_final)""" - return self.from_pulsar(obj), True \ No newline at end of file + # Check end_of_stream field to determine if this is the final message + is_final = getattr(obj, 'end_of_stream', True) + return self.from_pulsar(obj), is_final \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/services/llm.py b/trustgraph-base/trustgraph/schema/services/llm.py index 4665bc8a..3fd21937 100644 --- a/trustgraph-base/trustgraph/schema/services/llm.py +++ b/trustgraph-base/trustgraph/schema/services/llm.py @@ -1,5 +1,5 @@ -from pulsar.schema import Record, String, Array, Double, Integer +from pulsar.schema import Record, String, Array, Double, Integer, Boolean from ..core.topic import topic from ..core.primitives import Error @@ -11,6 +11,7 @@ from ..core.primitives import Error class TextCompletionRequest(Record): system = String() prompt = String() + streaming = Boolean() # Default false for backward compatibility class TextCompletionResponse(Record): error = Error() @@ -18,6 +19,7 @@ class TextCompletionResponse(Record): in_token = Integer() out_token = Integer() model = String() + end_of_stream = Boolean() # Indicates final message in stream ############################################################################ diff --git a/trustgraph-base/trustgraph/schema/services/prompt.py b/trustgraph-base/trustgraph/schema/services/prompt.py index 2567f471..edb569c9 100644 --- a/trustgraph-base/trustgraph/schema/services/prompt.py +++ b/trustgraph-base/trustgraph/schema/services/prompt.py @@ -1,4 +1,4 @@ -from pulsar.schema import Record, String, Map +from pulsar.schema import Record, String, Map, Boolean from ..core.primitives import Error from ..core.topic import topic @@ -24,6 +24,9 @@ class PromptRequest(Record): # JSON encoded values terms = Map(String()) + # Streaming support (default false for backward compatibility) + streaming = Boolean() + class PromptResponse(Record): # Error case @@ -35,4 +38,7 @@ class PromptResponse(Record): # JSON encoded object = String() + # Indicates final message in stream + end_of_stream = Boolean() + ############################################################################ \ No newline at end of file diff --git a/trustgraph-cli/trustgraph/cli/invoke_llm.py b/trustgraph-cli/trustgraph/cli/invoke_llm.py index d29286fb..da69dcd6 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_llm.py +++ b/trustgraph-cli/trustgraph/cli/invoke_llm.py @@ -6,17 +6,63 @@ and user prompt. Both arguments are required. import argparse import os import json -from trustgraph.api import Api +import uuid +import asyncio +from websockets.asyncio.client import connect -default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/') -def query(url, flow_id, system, prompt): +async def query(url, flow_id, system, prompt, streaming=True): - api = Api(url).flow().id(flow_id) + if not url.endswith("/"): + url += "/" - resp = api.text_completion(system=system, prompt=prompt) + url = url + "api/v1/socket" - print(resp) + mid = str(uuid.uuid4()) + + async with connect(url) as ws: + + req = { + "id": mid, + "service": "text-completion", + "flow": flow_id, + "request": { + "system": system, + "prompt": prompt, + "streaming": streaming + } + } + + await ws.send(json.dumps(req)) + + while True: + + msg = await ws.recv() + + obj = json.loads(msg) + + if "error" in obj: + raise RuntimeError(obj["error"]) + + if obj["id"] != mid: + continue + + if "response" in obj["response"]: + if streaming: + # Stream output to stdout without newline + print(obj["response"]["response"], end="", flush=True) + else: + # Non-streaming: print complete response + print(obj["response"]["response"]) + + if obj["complete"]: + if streaming: + # Add final newline after streaming + print() + break + + await ws.close() def main(): @@ -49,16 +95,23 @@ def main(): help=f'Flow ID (default: default)' ) + parser.add_argument( + '--no-streaming', + action='store_true', + help='Disable streaming (default: streaming enabled)' + ) + args = parser.parse_args() try: - query( + asyncio.run(query( url=args.url, - flow_id = args.flow_id, + flow_id=args.flow_id, system=args.system[0], prompt=args.prompt[0], - ) + streaming=not args.no_streaming + )) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/invoke_prompt.py b/trustgraph-cli/trustgraph/cli/invoke_prompt.py index 630a9281..c996c57d 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_prompt.py +++ b/trustgraph-cli/trustgraph/cli/invoke_prompt.py @@ -10,20 +10,76 @@ using key=value arguments on the command line, and these replace import argparse import os import json -from trustgraph.api import Api +import uuid +import asyncio +from websockets.asyncio.client import connect -default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/') -def query(url, flow_id, template_id, variables): +async def query(url, flow_id, template_id, variables, streaming=True): - api = Api(url).flow().id(flow_id) + if not url.endswith("/"): + url += "/" - resp = api.prompt(id=template_id, variables=variables) + url = url + "api/v1/socket" - if isinstance(resp, str): - print(resp) - else: - print(json.dumps(resp, indent=4)) + mid = str(uuid.uuid4()) + + async with connect(url) as ws: + + req = { + "id": mid, + "service": "prompt", + "flow": flow_id, + "request": { + "id": template_id, + "variables": variables, + "streaming": streaming + } + } + + await ws.send(json.dumps(req)) + + full_response = {"text": "", "object": ""} + + while True: + + msg = await ws.recv() + + obj = json.loads(msg) + + if "error" in obj: + raise RuntimeError(obj["error"]) + + if obj["id"] != mid: + continue + + response = obj["response"] + + # Handle text responses (streaming) + if "text" in response and response["text"]: + if streaming: + # Stream output to stdout without newline + print(response["text"], end="", flush=True) + full_response["text"] += response["text"] + else: + # Non-streaming: print complete response + print(response["text"]) + + # Handle object responses (JSON, never streamed) + if "object" in response and response["object"]: + full_response["object"] = response["object"] + + if obj["complete"]: + if streaming and full_response["text"]: + # Add final newline after streaming text + print() + elif full_response["object"]: + # Print JSON object (pretty-printed) + print(json.dumps(json.loads(full_response["object"]), indent=4)) + break + + await ws.close() def main(): @@ -59,6 +115,12 @@ def main(): specified multiple times''', ) + parser.add_argument( + '--no-streaming', + action='store_true', + help='Disable streaming (default: streaming enabled for text responses)' + ) + args = parser.parse_args() variables = {} @@ -73,12 +135,13 @@ specified multiple times''', try: - query( + asyncio.run(query( url=args.url, flow_id=args.flow_id, template_id=args.id[0], variables=variables, - ) + streaming=not args.no_streaming + )) except Exception as e: diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index eb79e472..5cf17b4d 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -32,7 +32,7 @@ from vertexai.generative_models import ( from anthropic import AnthropicVertex, RateLimitError from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk # Module logger logger = logging.getLogger(__name__) @@ -239,6 +239,123 @@ class Processor(LlmService): logger.error(f"VertexAI LLM exception: {e}", exc_info=True) raise e + def supports_streaming(self): + """VertexAI supports streaming for both Gemini and Claude models""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """ + Stream content generation from VertexAI (Gemini or Claude). + Yields LlmChunk objects with is_final=True on the last chunk. + """ + # Use provided model or fall back to default + model_name = model or self.default_model + # Use provided temperature or fall back to default + effective_temperature = temperature if temperature is not None else self.temperature + + logger.debug(f"Using model (streaming): {model_name}") + logger.debug(f"Using temperature: {effective_temperature}") + + try: + if 'claude' in model_name.lower(): + # Claude/Anthropic streaming + logger.debug(f"Streaming request to Anthropic model '{model_name}'...") + client = self._get_anthropic_client() + + total_in_tokens = 0 + total_out_tokens = 0 + + with client.messages.stream( + model=model_name, + system=system, + messages=[{"role": "user", "content": prompt}], + max_tokens=self.api_params['max_output_tokens'], + temperature=effective_temperature, + top_p=self.api_params['top_p'], + top_k=self.api_params['top_k'], + ) as stream: + # Stream text chunks + for text in stream.text_stream: + yield LlmChunk( + text=text, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + + # Get final message with token counts + final_message = stream.get_final_message() + total_in_tokens = final_message.usage.input_tokens + total_out_tokens = final_message.usage.output_tokens + + # Send final chunk with token counts + yield LlmChunk( + text="", + in_token=total_in_tokens, + out_token=total_out_tokens, + model=model_name, + is_final=True + ) + + logger.info(f"Input Tokens: {total_in_tokens}") + logger.info(f"Output Tokens: {total_out_tokens}") + + else: + # Gemini streaming + logger.debug(f"Streaming request to Gemini model '{model_name}'...") + full_prompt = system + "\n\n" + prompt + + llm, generation_config = self._get_gemini_model(model_name, effective_temperature) + + response = llm.generate_content( + full_prompt, + generation_config=generation_config, + safety_settings=self.safety_settings, + stream=True # Enable streaming + ) + + total_in_tokens = 0 + total_out_tokens = 0 + + # Stream chunks + for chunk in response: + if chunk.text: + yield LlmChunk( + text=chunk.text, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + + # Accumulate token counts if available + if hasattr(chunk, 'usage_metadata') and chunk.usage_metadata: + if hasattr(chunk.usage_metadata, 'prompt_token_count'): + total_in_tokens = chunk.usage_metadata.prompt_token_count + if hasattr(chunk.usage_metadata, 'candidates_token_count'): + total_out_tokens = chunk.usage_metadata.candidates_token_count + + # Send final chunk with token counts + yield LlmChunk( + text="", + in_token=total_in_tokens, + out_token=total_out_tokens, + model=model_name, + is_final=True + ) + + logger.info(f"Input Tokens: {total_in_tokens}") + logger.info(f"Output Tokens: {total_out_tokens}") + + except (google.api_core.exceptions.ResourceExhausted, RateLimitError) as e: + logger.warning(f"Hit rate limit during streaming: {e}") + raise TooManyRequests() + + except Exception as e: + logger.error(f"VertexAI streaming exception: {e}", exc_info=True) + raise e + @staticmethod def add_args(parser):