Scott Cogan
Add custom OpenAIModel implementation
74e6d68
raw
history blame contribute delete
594 Bytes
from openai import OpenAI
class OpenAIModel:
def __init__(self, model="gpt-4-turbo-preview", max_tokens=4096, temperature=0.7, api_key=None):
self.model = model
self.max_tokens = max_tokens
self.temperature = temperature
self.client = OpenAI(api_key=api_key)
def generate(self, messages):
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
max_tokens=self.max_tokens,
temperature=self.temperature
)
return response.choices[0].message.content