import time
import openai
from typing import Union, List, Dict
[docs]def get_gpt_output(message: Union[str, List[Dict[str, str]]],
model: str = "gpt-4-1106-preview",
max_tokens: int = 2048,
temperature: float = 1.0,
max_retry: int = 1,
sleep_time: int = 60,
json_object: bool = False) -> str:
"""
Call the OpenAI API to get the GPT model output for a given prompt.
Args:
message (Union[str, List[Dict[str, str]]]): The input message or a list of message dicts.
model (str): The model to use for completion. Default is "gpt-4-1106-preview".
max_tokens (int): Maximum number of tokens to generate. Default is 2048.
temperature (float): Sampling temperature. Default is 1.0.
max_retry (int): Maximum number of retries in case of an error. Default is 1.
sleep_time (int): Sleep time between retries in seconds. Default is 60.
json_object (bool): Whether to output in JSON format. Default is False.
Returns:
str: The completed text generated by the GPT model.
Raises:
Exception: If the completion fails after the maximum number of retries.
"""
if json_object:
if isinstance(message, str) and 'json' not in message.lower():
message = 'You are a helpful assistant designed to output JSON. ' + message
if isinstance(message, str):
messages = [{"role": "user", "content": message}]
else:
messages = message
kwargs = {"response_format": {"type": "json_object"}} if json_object else {}
for cnt in range(max_retry):
try:
chat = openai.OpenAI().chat.completions.create(
messages=messages,
model=model,
temperature=temperature,
max_tokens=max_tokens,
**kwargs
)
return chat.choices[0].message.content
except Exception as e:
print(f"Attempt {cnt} failed: {e}. Retrying after {sleep_time} seconds...")
time.sleep(sleep_time)
raise Exception("Failed to get GPT output after maximum retries")