diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index dbe6f54c..c172c5ea 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -11,7 +11,7 @@ import enum import logging from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk # Module logger logger = logging.getLogger(__name__) @@ -52,6 +52,8 @@ class ModelHandler: raise RuntimeError("format_request not implemented") def decode_response(self, response): raise RuntimeError("format_request not implemented") + def decode_stream_chunk(self, chunk): + raise RuntimeError("decode_stream_chunk not implemented") class Mistral(ModelHandler): def __init__(self): @@ -68,6 +70,11 @@ class Mistral(ModelHandler): def decode_response(self, response): response_body = json.loads(response.get("body").read()) return response_body['outputs'][0]['text'] + def decode_stream_chunk(self, chunk): + chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode()) + if 'outputs' in chunk_obj and len(chunk_obj['outputs']) > 0: + return chunk_obj['outputs'][0].get('text', '') + return '' # Llama 3 class Meta(ModelHandler): @@ -83,6 +90,9 @@ class Meta(ModelHandler): def decode_response(self, response): model_response = json.loads(response["body"].read()) return model_response["generation"] + def decode_stream_chunk(self, chunk): + chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode()) + return chunk_obj.get('generation', '') class Anthropic(ModelHandler): def __init__(self): @@ -108,6 +118,12 @@ class Anthropic(ModelHandler): def decode_response(self, response): model_response = json.loads(response["body"].read()) return model_response['content'][0]['text'] + def decode_stream_chunk(self, chunk): + chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode()) + if chunk_obj.get('type') == 'content_block_delta': + if 'delta' in chunk_obj and 'text' in chunk_obj['delta']: + return chunk_obj['delta']['text'] + return '' class Ai21(ModelHandler): def __init__(self): @@ -129,6 +145,12 @@ class Ai21(ModelHandler): content_str = content.decode('utf-8') content_json = json.loads(content_str) return content_json['choices'][0]['message']['content'] + def decode_stream_chunk(self, chunk): + chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode()) + if 'choices' in chunk_obj and len(chunk_obj['choices']) > 0: + delta = chunk_obj['choices'][0].get('delta', {}) + return delta.get('content', '') + return '' class Cohere(ModelHandler): def encode_request(self, system, prompt): @@ -142,6 +164,9 @@ class Cohere(ModelHandler): content_str = content.decode('utf-8') content_json = json.loads(content_str) return content_json['text'] + def decode_stream_chunk(self, chunk): + chunk_obj = json.loads(chunk.get('chunk').get('bytes').decode()) + return chunk_obj.get('text', '') Default=Mistral @@ -309,6 +334,78 @@ class Processor(LlmService): logger.error(f"Bedrock LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """Bedrock supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from Bedrock""" + model_name = model or self.default_model + 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: + variant = self._get_or_create_variant(model_name, effective_temperature) + promptbody = variant.encode_request(system, prompt) + + accept = 'application/json' + contentType = 'application/json' + + response = self.bedrock.invoke_model_with_response_stream( + body=promptbody, + modelId=model_name, + accept=accept, + contentType=contentType + ) + + total_input_tokens = 0 + total_output_tokens = 0 + + stream = response.get('body') + if stream: + for event in stream: + chunk = event.get('chunk') + if chunk: + # Decode the chunk text + text = variant.decode_stream_chunk(event) + if text: + yield LlmChunk( + text=text, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + + # Try to extract metadata from the event + metadata = event.get('metadata') + if metadata: + usage = metadata.get('usage') + if usage: + total_input_tokens = usage.get('inputTokens', 0) + total_output_tokens = usage.get('outputTokens', 0) + + # Send final chunk with token counts + yield LlmChunk( + text="", + in_token=total_input_tokens, + out_token=total_output_tokens, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except self.bedrock.exceptions.ThrottlingException as e: + logger.warning(f"Hit rate limit during streaming: {e}") + raise TooManyRequests() + + except Exception as e: + logger.error(f"Bedrock streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index d2d6f1ad..614c1362 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -11,7 +11,7 @@ import os import logging from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk # Module logger logger = logging.getLogger(__name__) @@ -55,7 +55,7 @@ class Processor(LlmService): self.max_output = max_output self.default_model = model - def build_prompt(self, system, content, temperature=None): + def build_prompt(self, system, content, temperature=None, stream=False): # Use provided temperature or fall back to default effective_temperature = temperature if temperature is not None else self.temperature @@ -73,6 +73,9 @@ class Processor(LlmService): "top_p": 1 } + if stream: + data["stream"] = True + body = json.dumps(data) return body @@ -157,6 +160,84 @@ class Processor(LlmService): logger.debug("Azure LLM processing complete") + def supports_streaming(self): + """Azure serverless endpoints support streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from Azure serverless endpoint""" + model_name = model or self.default_model + 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: + body = self.build_prompt(system, prompt, effective_temperature, stream=True) + + url = self.endpoint + api_key = self.token + + headers = { + 'Content-Type': 'application/json', + 'Authorization': f'Bearer {api_key}' + } + + response = requests.post(url, data=body, headers=headers, stream=True) + + if response.status_code == 429: + raise TooManyRequests() + + if response.status_code != 200: + raise RuntimeError("LLM failure") + + # Parse SSE stream + for line in response.iter_lines(): + if line: + line = line.decode('utf-8').strip() + if line.startswith('data: '): + data = line[6:] # Remove 'data: ' prefix + + if data == '[DONE]': + break + + try: + chunk_data = json.loads(data) + + if 'choices' in chunk_data and len(chunk_data['choices']) > 0: + delta = chunk_data['choices'][0].get('delta', {}) + content = delta.get('content') + if content: + yield LlmChunk( + text=content, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + except json.JSONDecodeError: + logger.warning(f"Failed to parse chunk: {data}") + continue + + # Send final chunk + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except TooManyRequests: + logger.warning("Rate limit exceeded during streaming") + raise TooManyRequests() + + except Exception as e: + logger.error(f"Azure streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index 5fecd3ac..2e7573d0 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -9,7 +9,7 @@ import os import logging from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk # Module logger logger = logging.getLogger(__name__) @@ -106,6 +106,65 @@ class Processor(LlmService): logger.error(f"Claude LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """Claude/Anthropic supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from Claude""" + model_name = model or self.default_model + 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: + with self.claude.messages.stream( + model=model_name, + max_tokens=self.max_output, + temperature=effective_temperature, + system=system, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt + } + ] + } + ] + ) as stream: + 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 for token counts + final_message = stream.get_final_message() + yield LlmChunk( + text="", + in_token=final_message.usage.input_tokens, + out_token=final_message.usage.output_tokens, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except anthropic.RateLimitError: + logger.warning("Rate limit exceeded during streaming") + raise TooManyRequests() + + except Exception as e: + logger.error(f"Claude streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index 9aebaa73..5093e556 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -13,7 +13,7 @@ import logging logger = logging.getLogger(__name__) from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk default_ident = "text-completion" @@ -98,6 +98,68 @@ class Processor(LlmService): logger.error(f"Cohere LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """Cohere supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from Cohere""" + model_name = model or self.default_model + 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: + stream = self.cohere.chat_stream( + model=model_name, + message=prompt, + preamble=system, + temperature=effective_temperature, + chat_history=[], + prompt_truncation='auto', + connectors=[] + ) + + total_input_tokens = 0 + total_output_tokens = 0 + + for event in stream: + if event.event_type == "text-generation": + if hasattr(event, 'text') and event.text: + yield LlmChunk( + text=event.text, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + elif event.event_type == "stream-end": + # Extract token counts from final event + if hasattr(event, 'response') and hasattr(event.response, 'meta'): + if hasattr(event.response.meta, 'billed_units'): + total_input_tokens = int(event.response.meta.billed_units.input_tokens) + total_output_tokens = int(event.response.meta.billed_units.output_tokens) + + # Send final chunk with token counts + yield LlmChunk( + text="", + in_token=total_input_tokens, + out_token=total_output_tokens, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except cohere.TooManyRequestsError: + logger.warning("Rate limit exceeded during streaming") + raise TooManyRequests() + + except Exception as e: + logger.error(f"Cohere streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index b9abcefa..1e9160ed 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -23,7 +23,7 @@ import logging logger = logging.getLogger(__name__) from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk default_ident = "text-completion" @@ -159,6 +159,67 @@ class Processor(LlmService): logger.error(f"GoogleAIStudio LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """Google AI Studio supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from Google AI Studio""" + model_name = model or self.default_model + 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}") + + generation_config = self._get_or_create_config(model_name, effective_temperature) + generation_config.system_instruction = system + + try: + response = self.client.models.generate_content_stream( + model=model_name, + config=generation_config, + contents=prompt, + ) + + total_input_tokens = 0 + total_output_tokens = 0 + + for chunk in response: + if hasattr(chunk, 'text') and 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'): + if hasattr(chunk.usage_metadata, 'prompt_token_count'): + total_input_tokens = int(chunk.usage_metadata.prompt_token_count) + if hasattr(chunk.usage_metadata, 'candidates_token_count'): + total_output_tokens = int(chunk.usage_metadata.candidates_token_count) + + # Send final chunk with token counts + yield LlmChunk( + text="", + in_token=total_input_tokens, + out_token=total_output_tokens, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except ResourceExhausted: + logger.warning("Rate limit exceeded during streaming") + raise TooManyRequests() + + except Exception as e: + logger.error(f"GoogleAIStudio streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py index 2b343583..801ed067 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py @@ -12,7 +12,7 @@ import logging logger = logging.getLogger(__name__) from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk default_ident = "text-completion" @@ -102,6 +102,57 @@ class Processor(LlmService): logger.error(f"Llamafile LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """LlamaFile supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from LlamaFile""" + model_name = model or self.default_model + 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}") + + prompt = system + "\n\n" + prompt + + try: + response = self.openai.chat.completions.create( + model=model_name, + messages=[{"role": "user", "content": prompt}], + temperature=effective_temperature, + max_tokens=self.max_output, + top_p=1, + frequency_penalty=0, + presence_penalty=0, + response_format={"type": "text"}, + stream=True + ) + + for chunk in response: + if chunk.choices and chunk.choices[0].delta.content: + yield LlmChunk( + text=chunk.choices[0].delta.content, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except Exception as e: + logger.error(f"LlamaFile streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py index 9e8e93dc..555d5c94 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py @@ -106,6 +106,57 @@ class Processor(LlmService): logger.error(f"LMStudio LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """LM Studio supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from LM Studio""" + model_name = model or self.default_model + 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}") + + prompt = system + "\n\n" + prompt + + try: + response = self.openai.chat.completions.create( + model=model_name, + messages=[{"role": "user", "content": prompt}], + temperature=effective_temperature, + max_tokens=self.max_output, + top_p=1, + frequency_penalty=0, + presence_penalty=0, + response_format={"type": "text"}, + stream=True + ) + + for chunk in response: + if chunk.choices and chunk.choices[0].delta.content: + yield LlmChunk( + text=chunk.choices[0].delta.content, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except Exception as e: + logger.error(f"LMStudio streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py index bcd00a0c..7952b1df 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py @@ -12,7 +12,7 @@ import logging logger = logging.getLogger(__name__) from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk default_ident = "text-completion" @@ -120,6 +120,67 @@ class Processor(LlmService): logger.error(f"Mistral LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """Mistral supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from Mistral""" + model_name = model or self.default_model + 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}") + + prompt = system + "\n\n" + prompt + + try: + stream = self.mistral.chat.stream( + model=model_name, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt + } + ] + } + ], + temperature=effective_temperature, + max_tokens=self.max_output, + top_p=1, + frequency_penalty=0, + presence_penalty=0, + response_format={"type": "text"} + ) + + for chunk in stream: + if chunk.data.choices and chunk.data.choices[0].delta.content: + yield LlmChunk( + text=chunk.data.choices[0].delta.content, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + + # Send final chunk + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except Exception as e: + logger.error(f"Mistral streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index db9586ea..3616e428 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -12,7 +12,7 @@ import logging logger = logging.getLogger(__name__) from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk default_ident = "text-completion" @@ -79,6 +79,62 @@ class Processor(LlmService): logger.error(f"Ollama LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """Ollama supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from Ollama""" + model_name = model or self.default_model + 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}") + + prompt = system + "\n\n" + prompt + + try: + stream = self.llm.generate( + model_name, + prompt, + options={'temperature': effective_temperature}, + stream=True + ) + + total_input_tokens = 0 + total_output_tokens = 0 + + for chunk in stream: + if 'response' in chunk and chunk['response']: + yield LlmChunk( + text=chunk['response'], + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + + # Accumulate token counts if available + if 'prompt_eval_count' in chunk: + total_input_tokens = int(chunk['prompt_eval_count']) + if 'eval_count' in chunk: + total_output_tokens = int(chunk['eval_count']) + + # Send final chunk with token counts + yield LlmChunk( + text="", + in_token=total_input_tokens, + out_token=total_output_tokens, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except Exception as e: + logger.error(f"Ollama streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py b/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py index 9bfc58be..63f8dbc4 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py @@ -12,7 +12,7 @@ import logging logger = logging.getLogger(__name__) from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk default_ident = "text-completion" @@ -121,6 +121,100 @@ class Processor(LlmService): logger.error(f"TGI LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """TGI supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from TGI""" + model_name = model or self.default_model + 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}") + + headers = { + "Content-Type": "application/json", + } + + request = { + "model": model_name, + "messages": [ + { + "role": "system", + "content": system, + }, + { + "role": "user", + "content": prompt, + } + ], + "max_tokens": self.max_output, + "temperature": effective_temperature, + "stream": True, + } + + try: + url = f"{self.base_url}/chat/completions" + + async with self.session.post( + url, + headers=headers, + json=request, + ) as response: + + if response.status != 200: + raise RuntimeError("Bad status: " + str(response.status)) + + # Parse SSE stream + async for line in response.content: + line = line.decode('utf-8').strip() + + if not line: + continue + + if line.startswith('data: '): + data = line[6:] # Remove 'data: ' prefix + + if data == '[DONE]': + break + + try: + import json + chunk_data = json.loads(data) + + # Extract text from chunk + if 'choices' in chunk_data and len(chunk_data['choices']) > 0: + choice = chunk_data['choices'][0] + if 'delta' in choice and 'content' in choice['delta']: + content = choice['delta']['content'] + if content: + yield LlmChunk( + text=content, + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + except json.JSONDecodeError: + logger.warning(f"Failed to parse chunk: {data}") + continue + + # Send final chunk + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except Exception as e: + logger.error(f"TGI streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py b/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py index 1cf7df49..af27830c 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py @@ -12,7 +12,7 @@ import logging logger = logging.getLogger(__name__) from .... exceptions import TooManyRequests -from .... base import LlmService, LlmResult +from .... base import LlmService, LlmResult, LlmChunk default_ident = "text-completion" @@ -113,6 +113,89 @@ class Processor(LlmService): logger.error(f"vLLM LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """vLLM supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """Stream content generation from vLLM""" + model_name = model or self.default_model + 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}") + + headers = { + "Content-Type": "application/json", + } + + request = { + "model": model_name, + "prompt": system + "\n\n" + prompt, + "max_tokens": self.max_output, + "temperature": effective_temperature, + "stream": True, + } + + try: + url = f"{self.base_url}/completions" + + async with self.session.post( + url, + headers=headers, + json=request, + ) as response: + + if response.status != 200: + raise RuntimeError("Bad status: " + str(response.status)) + + # Parse SSE stream + async for line in response.content: + line = line.decode('utf-8').strip() + + if not line: + continue + + if line.startswith('data: '): + data = line[6:] # Remove 'data: ' prefix + + if data == '[DONE]': + break + + try: + import json + chunk_data = json.loads(data) + + # Extract text from chunk + if 'choices' in chunk_data and len(chunk_data['choices']) > 0: + choice = chunk_data['choices'][0] + if 'text' in choice and choice['text']: + yield LlmChunk( + text=choice['text'], + in_token=None, + out_token=None, + model=model_name, + is_final=False + ) + except json.JSONDecodeError: + logger.warning(f"Failed to parse chunk: {data}") + continue + + # Send final chunk + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except Exception as e: + logger.error(f"vLLM streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser):