async def acompletion_with_retry(
    aclient: Any,
    max_retries: int,
    chat: bool = False,
    **kwargs: Any,
) -> Any:
    """Use tenacity to retry the async completion call."""
    retry_decorator = _create_retry_decorator(max_retries=max_retries)

    @retry_decorator
    async def _completion_with_retry(**kwargs: Any) -> Any:
        is_stream = kwargs.pop("stream", False)
        if chat:
            if is_stream:
                return await aclient.chat_stream(**kwargs)
            else:
                return await aclient.chat(**kwargs)
        else:
            if is_stream:
                return await aclient.generate_stream(**kwargs)
            else:
                return await aclient.generate(**kwargs)

    return await _completion_with_retry(**kwargs)
Use tenacity to retry the async completion call.
llama-index-integrations/llms/llama-index-llms-cohere/llama_index/llms/cohere/utils.py