itsprasun/pythonic-rag-fastapi-react
0
1from openai import OpenAI, AsyncOpenAI2from dotenv import load_dotenv3import os4 5load_dotenv()6 7 8class ChatOpenAI:9 def __init__(self, model_name: str = "gpt-4o-mini"):10 self.model_name = model_name11 self.openai_api_key = os.getenv("OPENAI_API_KEY")12 if self.openai_api_key is None:13 raise ValueError("OPENAI_API_KEY is not set")14 15 def run(self, messages, text_only: bool = True, **kwargs):16 if not isinstance(messages, list):17 raise ValueError("messages must be a list")18 19 client = OpenAI()20 response = client.chat.completions.create(21 model=self.model_name, messages=messages, **kwargs22 )23 24 if text_only:25 return response.choices[0].message.content26 27 return response28 29 async def astream(self, messages, **kwargs):30 if not isinstance(messages, list):31 raise ValueError("messages must be a list")32 33 client = AsyncOpenAI()34 35 stream = await client.chat.completions.create(36 model=self.model_name,37 messages=messages,38 stream=True,39 **kwargs40 )41 42 async for chunk in stream:43 content = chunk.choices[0].delta.content44 if content is not None:45 yield content46 