More model-specific support

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):
 

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"
 

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):