| import time | |
| import openai | |
| from typing import Union, List, Dict | |
| 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") | |