Spaces:
Build error
Build error
| 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 |