aimakerspace/openai_utils/chatmodel.py

import os
import openai
from typing import List, Dict, Union

class ChatOpenAI:
    def __init__(self, model_name: str = "gpt-4"):
        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: List[Dict[str, str]], text_only: bool = True) -> Union[str, Dict]:
        if not isinstance(messages, list):
            raise ValueError("messages must be a list")

        openai.api_key = self.openai_api_key
        response = openai.ChatCompletion.create(
            model=self.model_name,
            messages=messages
        )

        if text_only:
            return response.choices[0].message.content

        return response