text streaming ability added to OpenAI provider

This commit is contained in:
Barış Özmen
2024-11-02 22:06:20 +03:00
parent 92819112bb
commit efe1a62d73
+20
View File
@@ -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