Spaces:
No application file
No application file
File size: 1,208 Bytes
2be9eb9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
from typing import AsyncGenerator
import openai
from settings import Settings
class ChatOpenAI:
def __init__(self, settings: Settings) -> None:
self.model_name = settings.llm_model
self.client = openai.OpenAI(api_key=settings.openai_api_key)
self.aclient = openai.AsyncOpenAI(api_key=settings.openai_api_key)
def run(self, messages, text_only: bool = True):
if not isinstance(messages, list):
raise ValueError("messages must be a list")
chat_completion = self.client.chat.completions.create(
messages=messages, model=self.model_name
)
if text_only:
return chat_completion.choices[0].message.content
return chat_completion
async def astream(self, messages, **kwargs) -> AsyncGenerator:
if not isinstance(messages, list):
raise ValueError("messages must be a list")
stream = await self.aclient.chat.completions.create(
model=self.model_name, messages=messages, stream=True, **kwargs
)
async for chunk in stream:
content = chunk.choices[0].delta.content
if content is not None:
yield content
|