add max_attempts

src/gigachat.py

@@ -57,7 +57,7 @@ def check_auth_token(token_data: Dict[str, Any]) -> bool:
 
 token_data: Optional[Dict[str, Any]] = None
 
-def get_response(
+def get_response_json(
     prompt: str,
     model: str,
     timeout: int = 120,
@@ -65,6 +65,7 @@ def get_response(
     fuse_key_word: Optional[str] = None,
     use_giga_censor: bool = False,
     max_tokens: int = 128,
+    max_attempts: int = 5,
 ) -> requests.Response:
     """
     Send a text generation request to the API.
@@ -119,13 +120,24 @@ def get_response(
         'Accept': 'application/json',
         'Authorization': f'Bearer {token_data["access_token"]}'
     }
-    response = requests.post(url, headers=headers, data=payload, timeout=timeout)
-    return response
+
+    attempt_num = 0
+    while attempt_num < max_attempts:
+        try:
+            response = requests.post(url, headers=headers, data=payload, timeout=timeout)
+            response_dict = response.json()
+        except:
+            time.sleep(5)
+            attempt_num += 1
+            continue
+            
+    return response_dict
    
 def giga_generate(
     prompt: str, 
     model_version: str = "GigaChat-Max", 
-    max_tokens: int = 128
+    max_tokens: int = 128,
+    max_attempts: int = 5
 ) -> str:
     """
     Generate text using the GigaChat model.
@@ -138,17 +150,19 @@ def giga_generate(
     Returns:
         str: Generated text.
     """
-    response = get_response(
+    response_dict = get_response_json(
         prompt,
         model_version,
         use_giga_censor=False,
         max_tokens=max_tokens,
+        max_attempts=max_attempts,
     )
-    response_dict = response.json()
-
-    if response_dict['choices'][0]['finish_reason'] == 'blacklist':
-        print('GigaCensor triggered!')
-        return 'Censored Text'
-    else:
-        response_str = response_dict['choices'][0]['message']['content']
-        return response_str
+    try:
+        if response_dict['choices'][0]['finish_reason'] == 'blacklist':
+            print('GigaCensor triggered!')
+            return 'Censored Text'
+        else:
+            response_str = response_dict['choices'][0]['message']['content']
+            return response_str
+    except:
+        return prompt