feat(docvision): add retry mechanism to predict_layout

babeldoc/__init__.py

@@ -1 +1 @@
-__version__ = "0.1.14"
+__version__ = "0.1.15"

babeldoc/const.py

@@ -2,7 +2,7 @@
 import subprocess
 from pathlib import Path
 
-__version__ = "0.1.14"
+__version__ = "0.1.15"
 
 CACHE_FOLDER = Path.home() / ".cache" / "babeldoc"
 

babeldoc/docvision/rpc_doclayout.py

@@ -5,6 +5,10 @@
 import httpx
 import msgpack
 import numpy as np
+from tenacity import retry
+from tenacity import retry_if_exception_type
+from tenacity import stop_after_attempt
+from tenacity import wait_exponential
 
 from babeldoc.docvision.doclayout import DocLayoutModel
 from babeldoc.docvision.doclayout import YoloBox
@@ -34,6 +38,17 @@ def encode_image(image) -> bytes:
     return encoded
 
 
+@retry(
+    stop=stop_after_attempt(3),  # 最多重试3次
+    wait=wait_exponential(
+        multiplier=1, min=1, max=10
+    ),  # 指数退避策略，初始1秒，最大10秒
+    retry=retry_if_exception_type((httpx.HTTPError, Exception)),  # 针对哪些异常重试
+    before_sleep=lambda retry_state: logger.warning(
+        f"Request failed, retrying in {retry_state.next_action.sleep} seconds... "
+        f"(Attempt {retry_state.attempt_number}/3)"
+    ),
+)
 def predict_layout(
     image,
     host: str = "http://localhost:8000",
@@ -50,53 +65,49 @@ def predict_layout(
     Returns:
         List of predictions containing bounding boxes and classes
     """
-    try:
-        # Prepare request data
-        if not isinstance(image, list):
-            image = [image]
-        image_data = [encode_image(image) for image in image]
-        data = {
-            "image": image_data,
-            "imgsz": imgsz,
-        }
-
-        # Pack data using msgpack
-        packed_data = msgpack.packb(data, use_bin_type=True)
-        logger.debug(f"Packed data size: {len(packed_data)} bytes")
-
-        # Send request
-        logger.debug(f"Sending request to {host}/inference")
-        response = httpx.post(
-            f"{host}/inference",
-            data=packed_data,
-            headers={
-                "Content-Type": "application/msgpack",
-                "Accept": "application/msgpack",
-            },
-            timeout=300,
-            follow_redirects=True,
+    # Prepare request data
+    if not isinstance(image, list):
+        image = [image]
+    image_data = [encode_image(image) for image in image]
+    data = {
+        "image": image_data,
+        "imgsz": imgsz,
+    }
+
+    # Pack data using msgpack
+    packed_data = msgpack.packb(data, use_bin_type=True)
+    logger.debug(f"Packed data size: {len(packed_data)} bytes")
+
+    # Send request
+    logger.debug(f"Sending request to {host}/inference")
+    response = httpx.post(
+        f"{host}/inference",
+        data=packed_data,
+        headers={
+            "Content-Type": "application/msgpack",
+            "Accept": "application/msgpack",
+        },
+        timeout=300,
+        follow_redirects=True,
+    )
+
+    logger.debug(f"Response status: {response.status_code}")
+    logger.debug(f"Response headers: {response.headers}")
+
+    if response.status_code == 200:
+        try:
+            result = msgpack.unpackb(response.content, raw=False)
+            return result
+        except Exception as e:
+            logger.exception(f"Failed to unpack response: {e!s}")
+            raise
+    else:
+        logger.error(f"Request failed with status {response.status_code}")
+        logger.error(f"Response content: {response.content}")
+        raise Exception(
+            f"Request failed with status {response.status_code}: {response.text}",
         )
 
-        logger.debug(f"Response status: {response.status_code}")
-        logger.debug(f"Response headers: {response.headers}")
-
-        if response.status_code == 200:
-            try:
-                result = msgpack.unpackb(response.content, raw=False)
-                return result
-            except Exception as e:
-                logger.exception(f"Failed to unpack response: {e!s}")
-                raise
-        else:
-            logger.error(f"Request failed with status {response.status_code}")
-            logger.error(f"Response content: {response.content}")
-            raise Exception(
-                f"Request failed with status {response.status_code}: {response.text}",
-            )
-    except Exception as e:
-        logger.exception(f"Unexpected error: {e!s}")
-        raise
-
 
 class RpcDocLayoutModel(DocLayoutModel):
     """DocLayoutModel implementation that uses RPC service."""

babeldoc/main.py

@@ -23,7 +23,7 @@
 from babeldoc.translation_config import TranslationConfig
 
 logger = logging.getLogger(__name__)
-__version__ = "0.1.14"
+__version__ = "0.1.15"
 
 
 def create_parser():

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "BabelDOC"
-version = "0.1.14"
+version = "0.1.15"
 description = "Yet Another Document Translator"
 license = "AGPL-3.0"
 readme = "README.md"
@@ -30,6 +30,7 @@ dependencies = [
     "xsdata[cli,lxml,soap]>=24.12",
     "msgpack>=1.1.0",
     "pydantic>=2.10.6",
+    "tenacity>=9.0.0",
 ]
 
 [project.urls]
@@ -131,7 +132,7 @@ pythonpath = [".", "src"]
 testpaths = ["tests"]
 
 [bumpver]
-current_version = "0.1.14"
+current_version = "0.1.15"
 version_pattern = "MAJOR.MINOR.PATCH[.PYTAGNUM]"
 
 [bumpver.file_patterns]
@@ -147,4 +148,4 @@ version_pattern = "MAJOR.MINOR.PATCH[.PYTAGNUM]"
 ]
 "babeldoc/const.py" = [
     '__version__ = "{version}"'
-]
+]