| from openai import OpenAI | |
| from dotenv import load_dotenv | |
| import os | |
| load_dotenv() | |
| class ChatOpenAI: | |
| def __init__(self, model_name: str = "gpt-3.5-turbo"): | |
| 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): | |
| 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 | |
| ) | |
| if text_only: | |
| return response.choices[0].message.content | |
| return response | |