Skip to content

Commit

Permalink
feat: Update Max Files per Batch Request to 1000 (#91)
Browse files Browse the repository at this point in the history
* feat: Update Max Files per Batch Request to 1000
- https://cloud.google.com/document-ai/quotas#content_limits

* test: Update `test_create_batches_with_invalid_batch_size()`

* test: Add `batch_size` parameter to `test_create_batches_with_large_folder()`
  • Loading branch information
holtskinner committed Apr 3, 2023
1 parent e5ef356 commit 3bbc0f0
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 6 deletions.
7 changes: 4 additions & 3 deletions google/cloud/documentai_toolbox/constants.py
Expand Up @@ -21,16 +21,17 @@

FILE_CHECK_REGEX = r"(.*[.].*$)"

# https://cloud.google.com/document-ai/quotas
BATCH_MAX_FILES = 50
# https://cloud.google.com/document-ai/quotas#content_limits
BATCH_MAX_FILES = 1000
# 1GB in Bytes
BATCH_MAX_FILE_SIZE = 1073741824
BATCH_MAX_REQUESTS = 5

# https://cloud.google.com/document-ai/docs/file-types
VALID_MIME_TYPES = {
"application/pdf",
"image/bmp" "image/gif",
"image/bmp",
"image/gif",
"image/jpeg",
"image/png",
"image/tiff",
Expand Down
6 changes: 3 additions & 3 deletions tests/unit/test_utilities.py
Expand Up @@ -350,10 +350,10 @@ def test_create_batches_with_3_documents(mock_storage, capfd):
def test_create_batches_with_invalid_batch_size():
with pytest.raises(
ValueError,
match="Batch size must be less than 50. You provided 51.",
match="Batch size must be less than 1000. You provided 1001.",
):
utilities.create_batches(
gcs_bucket_name=test_bucket, gcs_prefix=test_prefix, batch_size=51
gcs_bucket_name=test_bucket, gcs_prefix=test_prefix, batch_size=1001
)


Expand All @@ -373,7 +373,7 @@ def test_create_batches_with_large_folder(mock_storage, capfd):
client.list_blobs.return_value = mock_blobs

actual = utilities.create_batches(
gcs_bucket_name=test_bucket, gcs_prefix=test_prefix
gcs_bucket_name=test_bucket, gcs_prefix=test_prefix, batch_size=50
)

mock_storage.Client.assert_called_once()
Expand Down

0 comments on commit 3bbc0f0

Please sign in to comment.