From efe1a62d7338f49f8ae2a5bdfc6468d757c3324c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bar=C4=B1=C5=9F=20=C3=96zmen?= <hbaristr@gmail.com>
Date: Sat, 2 Nov 2024 22:06:20 +0300
Subject: [PATCH] text streaming ability added to OpenAI provider

---
 simplemind/providers/openai.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/simplemind/providers/openai.py b/simplemind/providers/openai.py
index fb197e5..9779e01 100644
--- a/simplemind/providers/openai.py
+++ b/simplemind/providers/openai.py
@@ -106,3 +106,23 @@ class OpenAI(BaseProvider):
             **{**self.DEFAULT_KWARGS, **kwargs},
         )
         return response.choices[0].message.content
+    
+    @logger
+    def generate_stream_text(self, prompt: str, *, llm_model: str | None = None, **kwargs):
+        """Generate streaming text using the OpenAI API.
+        
+        Yields chunks of text as they are generated by the model.
+        """
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+        response = self.client.chat.completions.create(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            stream=True,  # Enable streaming
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+        
+        for chunk in response:
+            if chunk.choices[0].delta.content is not None:
+                yield chunk.choices[0].delta.content