From 1be0d12cf01f401811544b63001c17374207a058 Mon Sep 17 00:00:00 2001 From: Dev Aggarwal Date: Mon, 20 Nov 2023 20:17:03 +0530 Subject: [PATCH] retry google sheets rate limits --- recipes/DocExtract.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/recipes/DocExtract.py b/recipes/DocExtract.py index 182ce97c6..cca9848f7 100644 --- a/recipes/DocExtract.py +++ b/recipes/DocExtract.py @@ -4,8 +4,10 @@ import typing import requests +from aifail import retry_if from django.db.models import IntegerChoices from furl import furl +from googleapiclient.errors import HttpError from pydantic import BaseModel import gooey_ui as st @@ -20,7 +22,6 @@ run_google_translate, audio_url_to_wav, ) -from daras_ai_v2.glossary import glossary_input from daras_ai_v2.azure_doc_extract import azure_doc_extract_pages from daras_ai_v2.base import BasePage from daras_ai_v2.doc_search_settings_widgets import document_uploader @@ -31,6 +32,7 @@ flatapply_parallel, ) from daras_ai_v2.gdrive_downloader import is_gdrive_url, gdrive_download +from daras_ai_v2.glossary import glossary_input from daras_ai_v2.language_model import run_language_model, LargeLanguageModels from daras_ai_v2.language_model_settings_widgets import language_model_settings from daras_ai_v2.loom_video_widget import youtube_video @@ -431,6 +433,13 @@ def process_source( update_cell(spreadsheet_id, row, Columns.summary.value, summary) +def google_api_should_retry(e: Exception) -> bool: + return isinstance(e, HttpError) and ( + e.resp.status in (408, 429) or e.resp.status > 500 + ) + + +@retry_if(google_api_should_retry) def update_cell(spreadsheet_id: str, row: int, col: int, value: str): get_spreadsheet_service().values().update( spreadsheetId=spreadsheet_id,