From 9d234da5289e595b879eb8f57ffc00676190f6b4 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 25 Nov 2025 19:50:01 +0000 Subject: [PATCH] More model-specific support --- .../model/text_completion/azure_openai/llm.py | 71 ++++++++++++++++++- .../model/text_completion/lmstudio/llm.py | 2 +- .../model/text_completion/openai/llm.py | 71 ++++++++++++++++++- 3 files changed, 141 insertions(+), 3 deletions(-) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index 2442c283..950c006a 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -14,7 +14,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" @@ -125,6 +125,75 @@ class Processor(LlmService): logger.debug("Azure OpenAI LLM processing complete") + def supports_streaming(self): + """Azure OpenAI supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """ + Stream content generation from Azure OpenAI. + 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}") + + prompt = system + "\n\n" + prompt + + try: + response = self.openai.chat.completions.create( + model=model_name, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt + } + ] + } + ], + temperature=effective_temperature, + max_tokens=self.max_output, + top_p=1, + stream=True # Enable streaming + ) + + # Stream chunks + 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 + ) + + # Send final chunk + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except RateLimitError: + logger.warning("Rate limit exceeded during streaming") + raise TooManyRequests() + + except Exception as e: + logger.error(f"Azure OpenAI 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 a5464368..9e8e93dc 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/lmstudio/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" diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index d2698589..4da1378b 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/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__) @@ -118,6 +118,75 @@ class Processor(LlmService): logger.error(f"OpenAI LLM exception ({type(e).__name__}): {e}", exc_info=True) raise e + def supports_streaming(self): + """OpenAI supports streaming""" + return True + + async def generate_content_stream(self, system, prompt, model=None, temperature=None): + """ + Stream content generation from OpenAI. + 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}") + + prompt = system + "\n\n" + prompt + + try: + response = self.openai.chat.completions.create( + model=model_name, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt + } + ] + } + ], + temperature=effective_temperature, + max_tokens=self.max_output, + stream=True # Enable streaming + ) + + # Stream chunks + 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 + ) + + # Note: OpenAI doesn't provide token counts in streaming mode + # Send final chunk without token counts + yield LlmChunk( + text="", + in_token=None, + out_token=None, + model=model_name, + is_final=True + ) + + logger.debug("Streaming complete") + + except RateLimitError: + logger.warning("Hit rate limit during streaming") + raise TooManyRequests() + + except Exception as e: + logger.error(f"OpenAI streaming exception ({type(e).__name__}): {e}", exc_info=True) + raise e + @staticmethod def add_args(parser):