Team Ai
Apppublic

itsprasun/pythonic-rag-fastapi-react

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
chatmodel.py46 linesDownload Raw Back to openai_utils
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