fix: support dynamic timeouts for openai and azure

This commit is contained in:
Krrish Dholakia 2023-12-30 12:14:02 +05:30
parent 77be3e3114
commit c33c1d85bb
3 changed files with 36 additions and 24 deletions

View file

@ -306,7 +306,7 @@ class OpenAIChatCompletion(BaseLLM):
)
else:
openai_client = client
response = openai_client.chat.completions.create(**data) # type: ignore
response = openai_client.chat.completions.create(**data, timeout=timeout) # type: ignore
stringified_response = response.model_dump_json()
logging_obj.post_call(
input=messages,
@ -383,7 +383,9 @@ class OpenAIChatCompletion(BaseLLM):
},
)
response = await openai_aclient.chat.completions.create(**data)
response = await openai_aclient.chat.completions.create(
**data, timeout=timeout
)
stringified_response = response.model_dump_json()
logging_obj.post_call(
input=data["messages"],
@ -431,7 +433,7 @@ class OpenAIChatCompletion(BaseLLM):
"complete_input_dict": data,
},
)
response = openai_client.chat.completions.create(**data)
response = openai_client.chat.completions.create(**data, timeout=timeout)
streamwrapper = CustomStreamWrapper(
completion_stream=response,
model=model,
@ -476,7 +478,9 @@ class OpenAIChatCompletion(BaseLLM):
},
)
response = await openai_aclient.chat.completions.create(**data)
response = await openai_aclient.chat.completions.create(
**data, timeout=timeout
)
streamwrapper = CustomStreamWrapper(
completion_stream=response,
model=model,
@ -522,7 +526,7 @@ class OpenAIChatCompletion(BaseLLM):
)
else:
openai_aclient = client
response = await openai_aclient.embeddings.create(**data) # type: ignore
response = await openai_aclient.embeddings.create(**data, timeout=timeout) # type: ignore
stringified_response = response.model_dump_json()
## LOGGING
logging_obj.post_call(
@ -584,7 +588,7 @@ class OpenAIChatCompletion(BaseLLM):
openai_client = client
## COMPLETION CALL
response = openai_client.embeddings.create(**data) # type: ignore
response = openai_client.embeddings.create(**data, timeout=timeout) # type: ignore
## LOGGING
logging_obj.post_call(
input=input,
@ -629,7 +633,7 @@ class OpenAIChatCompletion(BaseLLM):
)
else:
openai_aclient = client
response = await openai_aclient.images.generate(**data) # type: ignore
response = await openai_aclient.images.generate(**data, timeout=timeout) # type: ignore
stringified_response = response.model_dump_json()
## LOGGING
logging_obj.post_call(
@ -669,9 +673,9 @@ class OpenAIChatCompletion(BaseLLM):
if not isinstance(max_retries, int):
raise OpenAIError(status_code=422, message="max retries must be an int")
# if aembedding == True:
# response = self.aembedding(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries) # type: ignore
# return response
if aimg_generation == True:
response = self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries) # type: ignore
return response
if client is None:
openai_client = OpenAI(
@ -697,7 +701,7 @@ class OpenAIChatCompletion(BaseLLM):
)
## COMPLETION CALL
response = openai_client.images.generate(**data) # type: ignore
response = openai_client.images.generate(**data, timeout=timeout) # type: ignore
## LOGGING
logging_obj.post_call(
input=input,