(fix) completion: max_retries using OpenAI client

This commit is contained in:
ishaan-jaff 2023-11-20 16:53:18 -08:00
parent 2bd934e56c
commit 11ec2710c6
4 changed files with 19 additions and 16 deletions

View file

@ -207,7 +207,7 @@ class OpenAIChatCompletion(BaseLLM):
elif optional_params.get("stream", False):
return self.streaming(logging_obj=logging_obj, data=data, model=model, api_base=api_base, api_key=api_key, timeout=timeout)
else:
openai_client = OpenAI(api_key=api_key, base_url=api_base, http_client=litellm.client_session, timeout=timeout)
openai_client = OpenAI(api_key=api_key, base_url=api_base, http_client=litellm.client_session, timeout=timeout, max_retries=data.pop("max_retries", 2))
response = openai_client.chat.completions.create(**data) # type: ignore
logging_obj.post_call(
input=None,
@ -249,7 +249,7 @@ class OpenAIChatCompletion(BaseLLM):
api_base: Optional[str]=None):
response = None
try:
openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base, http_client=litellm.aclient_session, timeout=timeout)
openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base, http_client=litellm.aclient_session, timeout=timeout, max_retries=data.pop("max_retries", 2))
response = await openai_aclient.chat.completions.create(**data)
return convert_to_model_response_object(response_object=json.loads(response.model_dump_json()), model_response_object=model_response)
except Exception as e:
@ -269,7 +269,7 @@ class OpenAIChatCompletion(BaseLLM):
api_key: Optional[str]=None,
api_base: Optional[str]=None
):
openai_client = OpenAI(api_key=api_key, base_url=api_base, http_client=litellm.client_session, timeout=timeout)
openai_client = OpenAI(api_key=api_key, base_url=api_base, http_client=litellm.client_session, timeout=timeout, max_retries=data.pop("max_retries", 2))
response = openai_client.chat.completions.create(**data)
streamwrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="openai",logging_obj=logging_obj)
for transformed_chunk in streamwrapper:
@ -284,7 +284,7 @@ class OpenAIChatCompletion(BaseLLM):
api_base: Optional[str]=None):
response = None
try:
openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base, http_client=litellm.aclient_session, timeout=timeout)
openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base, http_client=litellm.aclient_session, timeout=timeout, max_retries=data.pop("max_retries", 2))
response = await openai_aclient.chat.completions.create(**data)
streamwrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="openai",logging_obj=logging_obj)
async for transformed_chunk in streamwrapper:
@ -309,7 +309,7 @@ class OpenAIChatCompletion(BaseLLM):
super().embedding()
exception_mapping_worked = False
try:
openai_client = OpenAI(api_key=api_key, base_url=api_base, http_client=litellm.client_session)
openai_client = OpenAI(api_key=api_key, base_url=api_base, http_client=litellm.client_session, max_retries=data.pop("max_retries", 2))
model = model
data = {
"model": model,