File size: 1,308 Bytes
234eac0 |
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 40 41 42 43 44 45 46 |
from openai import OpenAI, AsyncOpenAI
from dotenv import load_dotenv
import os
load_dotenv()
class ChatOpenAI:
def __init__(self, model_name: str = "gpt-4o-mini"):
self.model_name = model_name
self.openai_api_key = os.getenv("OPENAI_API_KEY")
if self.openai_api_key is None:
raise ValueError("OPENAI_API_KEY is not set")
def run(self, messages, text_only: bool = True, **kwargs):
if not isinstance(messages, list):
raise ValueError("messages must be a list")
client = OpenAI()
response = client.chat.completions.create(
model=self.model_name, messages=messages, **kwargs
)
if text_only:
return response.choices[0].message.content
return response
async def astream(self, messages, **kwargs):
if not isinstance(messages, list):
raise ValueError("messages must be a list")
client = AsyncOpenAI()
stream = await client.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
|