From 05044dfaf2edda87abee168b69fe0215ebf75611 Mon Sep 17 00:00:00 2001
From: Yi Ding <cuauty@users.noreply.github.com>
Date: Wed, 20 Nov 2024 16:40:22 +0800
Subject: [PATCH] [Update] Support new error code for Bailing model (#1702)

* support new error code

* fix the lint problems
---
 opencompass/models/bailing_api_oc.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/opencompass/models/bailing_api_oc.py b/opencompass/models/bailing_api_oc.py
index 0b721bff..316f529b 100644
--- a/opencompass/models/bailing_api_oc.py
+++ b/opencompass/models/bailing_api_oc.py
@@ -2,6 +2,7 @@ import concurrent
 import concurrent.futures
 import os
 import socket
+import time
 import traceback
 from typing import Dict, List, Optional, Union
 
@@ -20,6 +21,8 @@ from .base_api import BaseAPIModel
 
 PromptType = Union[PromptList, str]
 
+BAILING_RETRY_DELAY: int = 30
+
 
 class HTTPAdapterWithSocketOptions(HTTPAdapter):
 
@@ -200,6 +203,9 @@ class BailingAPI(BaseAPIModel):
                     break  # success
                 elif response.status_code == 426:
                     retry_num += 1  # retry
+                elif response.status_code in [429, 500, 504]:
+                    time.sleep(BAILING_RETRY_DELAY)
+                    retry_num += 1  # retry
                 else:
                     raise ValueError(f'Status code = {response.status_code}')
             else: