Retry logic for AzureOpenAIEmbeddings

[jerichosiahaya]

I’m using AzureOpenAIEmbeddings and encountered an issue with the rate limit. I checked the code for OpenAIEmbeddings, which includes a retry logic function. However, when I checked AzureOpenAIEmbeddings, I noticed there is no retry function. I’m wondering why both don’t have retry logic, are there any other concerns?

Retry function for OpenAIEmbeddings:

langchain/libs/community/langchain_community/embeddings/openai.py

Lines 44 to 68 in 5a2cfb4

	def _create_retry_decorator(embeddings: OpenAIEmbeddings) -> Callable[[Any], Any]:
	import openai
	# Wait 2^x * 1 second between each retry starting with
	# retry_min_seconds seconds, then up to retry_max_seconds seconds,
	# then retry_max_seconds seconds afterwards
	# retry_min_seconds and retry_max_seconds are optional arguments of
	# OpenAIEmbeddings
	return retry(
	reraise=True,
	stop=stop_after_attempt(embeddings.max_retries),
	wait=wait_exponential(
	multiplier=1,
	min=embeddings.retry_min_seconds,
	max=embeddings.retry_max_seconds,
	),
	retry=(
	retry_if_exception_type(openai.error.Timeout) # type: ignore[attr-defined]
	| retry_if_exception_type(openai.error.APIError) # type: ignore[attr-defined]
	| retry_if_exception_type(openai.error.APIConnectionError) # type: ignore[attr-defined]
	| retry_if_exception_type(openai.error.RateLimitError) # type: ignore[attr-defined]
	| retry_if_exception_type(openai.error.ServiceUnavailableError) # type: ignore[attr-defined]
	),
	before_sleep=before_sleep_log(logger, logging.WARNING),
	)

langchain/libs/community/langchain_community/embeddings/openai.py

Lines 495 to 501 in 5a2cfb4

	batched_embeddings: List[List[float]] = []
	for i in _iter:
	response = embed_with_retry(
	self,
	input=tokens[i : i + _chunk_size],
	**self._invocation_params,
	)

Currently, my simple solution to solve this in AzureOpenAIEmbeddings is to get RateLimitError exception and then cooldown for 30 seconds before try for next attempt.