# the async version is adapted from https://gist.github.com/neubig/80de662fb3e225c18172ec218be4917a from __future__ import annotations import os import yaml import openai import ast import pdb import asyncio from typing import Any, List import os import pathlib import openai from openai import OpenAI, AsyncOpenAI import re class OpenAIChat(): def __init__( self, model_name='gpt-3.5-turbo', max_tokens=2500, temperature=0, top_p=1, request_timeout=120, ): if 'gpt' not in model_name: openai.api_base = "http://localhost:8000/v1" else: #openai.api_base = "https://api.openai.com/v1" openai.api_key = os.environ.get("OPENAI_API_KEY", None) assert openai.api_key is not None, "Please set the OPENAI_API_KEY environment variable." assert openai.api_key !='', "Please set the OPENAI_API_KEY environment variable." self.client = AsyncOpenAI() self.config = { 'model_name': model_name, 'max_tokens': max_tokens, 'temperature': temperature, 'top_p': top_p, 'request_timeout': request_timeout, } def extract_list_from_string(self, input_string): # pattern = r'\[.*\]' # result = re.search(pattern, input_string) # if result: # return result.group() # else: # return None start_index = input_string.find('[') end_index = input_string.rfind(']') if start_index != -1 and end_index != -1 and start_index < end_index: return input_string[start_index:end_index + 1] else: return None def extract_dict_from_string(self, input_string): start_index = input_string.find('{') end_index = input_string.rfind('}') if start_index != -1 and end_index != -1 and start_index < end_index: return input_string[start_index:end_index + 1] else: return None def _boolean_fix(self, output): return output.replace("true", "True").replace("false", "False") def _type_check(self, output, expected_type): try: output_eval = ast.literal_eval(output) if not isinstance(output_eval, expected_type): return None return output_eval except: ''' if(expected_type == List): valid_output = self.extract_list_from_string(output) output_eval = ast.literal_eval(valid_output) if not isinstance(output_eval, expected_type): return None return output_eval elif(expected_type == dict): valid_output = self.extract_dict_from_string(output) output_eval = ast.literal_eval(valid_output) if not isinstance(output_eval, expected_type): return None return output_eval ''' return None async def dispatch_openai_requests(self, messages_list,) -> list[str]: """ Dispatches requests to OpenAI API asynchronously. Args: messages_list: List of messages to be sent to OpenAI ChatCompletion API. Returns: List of responses from OpenAI API. """ async def _request_with_retry(messages, retry=3): for attempt in range(retry): try: response = await self.client.chat.completions.create( model=self.config['model_name'], messages=messages, max_tokens=self.config['max_tokens'], temperature=self.config['temperature'], top_p=self.config['top_p'] ) return response except openai.RateLimitError as e: await asyncio.sleep((2 ** attempt) * 0.5) # exponential backoff except (openai.Timeout, openai.APIError) as e: await asyncio.sleep((2 ** attempt) * 0.5) # exponential backoff except Exception as e: # Log unexpected exception for further investigation await asyncio.sleep((2 ** attempt) * 0.5) # fallback in case of unknown errors raise RuntimeError("All retries failed for OpenAI API request") async_responses = [ _request_with_retry(messages) for messages in messages_list ] return await asyncio.gather(*async_responses, return_exceptions=True) def run(self, messages_list, expected_type): retry = 1 responses = [None for _ in range(len(messages_list))] messages_list_cur_index = [i for i in range(len(messages_list))] while retry > 0 and len(messages_list_cur_index) > 0: messages_list_cur = [messages_list[i] for i in messages_list_cur_index] predictions = asyncio.run(self.dispatch_openai_requests( messages_list=messages_list_cur, )) preds = [self._type_check(self._boolean_fix(prediction.choices[0].message.content), expected_type) if prediction is not None else None for prediction in predictions] finised_index = [] for i, pred in enumerate(preds): if pred is not None: responses[messages_list_cur_index[i]] = pred finised_index.append(messages_list_cur_index[i]) messages_list_cur_index = [i for i in messages_list_cur_index if i not in finised_index] retry -= 1 return responses # class OpenAIEmbed(): # def __init__(): # openai.api_key = os.environ.get("OPENAI_API_KEY", None) # assert openai.api_key is not None, "Please set the OPENAI_API_KEY environment variable." # assert openai.api_key != '', "Please set the OPENAI_API_KEY environment variable." # async def create_embedding(self, text, retry=3): # for _ in range(retry): # try: # response = await openai.Embedding.acreate(input=text, model="text-embedding-ada-002") # return response # except openai.error.RateLimitError: # print('Rate limit error, waiting for 1 second...') # await asyncio.sleep(1) # except openai.error.APIError: # print('API error, waiting for 1 second...') # await asyncio.sleep(1) # except openai.error.Timeout: # print('Timeout error, waiting for 1 second...') # await asyncio.sleep(1) # return None # async def process_batch(self, batch, retry=3): # tasks = [self.create_embedding(text, retry=retry) for text in batch] # return await asyncio.gather(*tasks) # if __name__ == "__main__": # chat = OpenAIChat(model_name='llama-2-7b-chat-hf') # predictions = asyncio.run(chat.async_run( # messages_list=[ # [{"role": "user", "content": "show either 'ab' or '['a']'. Do not do anything else."}], # ] * 20, # expected_type=List, # )) # print(predictions) # Usage # embed = OpenAIEmbed() # batch = ["string1", "string2", "string3", "string4", "string5", "string6", "string7", "string8", "string9", "string10"] # Your batch of strings # embeddings = asyncio.run(embed.process_batch(batch, retry=3)) # for embedding in embeddings: # print(embedding["data"][0]["embedding"])