backend/utils/chat.py

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