Update app.py

app.py

@@ -2,6 +2,8 @@ import os
 import time
 import logging
 import requests
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
 import json
 import uuid
 import concurrent.futures
@@ -28,6 +30,24 @@ MODELS_ENDPOINT = "https://api-st.siliconflow.cn/v1/models"
 EMBEDDINGS_ENDPOINT = "https://api-st.siliconflow.cn/v1/embeddings"
 IMAGE_ENDPOINT = "https://api-st.siliconflow.cn/v1/images/generations"
 
+def requests_session_with_retries(
+    retries=3, backoff_factor=0.3, status_forcelist=(500, 502, 504)
+):
+    session = requests.Session()
+    retry = Retry(
+        total=retries,
+        read=retries,
+        connect=retries,
+        backoff_factor=backoff_factor,
+        status_forcelist=status_forcelist,
+    )
+    adapter = HTTPAdapter(max_retries=retry)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session
+
+session = requests_session_with_retries()
+
 app = Flask(__name__)
 app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1)
 
@@ -61,7 +81,7 @@ def get_credit_summary(api_key):
 
     for attempt in range(max_retries):
         try:
-            response = requests.get(API_ENDPOINT, headers=headers, timeout=2)
+            response = session.get(API_ENDPOINT, headers=headers, timeout=2)
             response.raise_for_status()
             data = response.json().get("data", {})
             total_balance = data.get("totalBalance", 0)
@@ -93,7 +113,7 @@ def test_model_availability(api_key, model_name):
         "Content-Type": "application/json"
     }
     try:
-        response = requests.post(
+        response = session.post(
             TEST_MODEL_ENDPOINT,
             headers=headers,
             json={
@@ -117,7 +137,7 @@ def test_model_availability(api_key, model_name):
     
 def create_base64_markdown_image(image_url):
     try:
-        response = requests.get(image_url, stream=True)
+        response = session.get(image_url, stream=True)
         response.raise_for_status()
 
         image_data = response.content
@@ -250,7 +270,7 @@ def test_embedding_model_availability(api_key, model_name):
         "Content-Type": "application/json"
     }
     try:
-        response = requests.post(
+        response = session.post(
             EMBEDDINGS_ENDPOINT,
             headers=headers,
             json={
@@ -354,7 +374,7 @@ def get_all_models(api_key, sub_type):
         "Content-Type": "application/json"
     }
     try:
-        response = requests.get(
+        response = session.get(
             MODELS_ENDPOINT,
             headers=headers,
             params={"sub_type": sub_type}
@@ -765,7 +785,7 @@ def handsome_images_generations():
                         print(f"image_url: {image_url}")
                         if data.get("response_format") == "b64_json":
                            try:
-                                image_data = requests.get(image_url, stream=True).raw
+                                image_data = session.get(image_url, stream=True).raw
                                 image = Image.open(image_data)
                                 buffered = io.BytesIO()
                                 image.save(buffered, format="PNG")
@@ -1363,4 +1383,4 @@ if __name__ == '__main__':
         debug=False,
         host='0.0.0.0',
         port=int(os.environ.get('PORT', 7860))
-    )
+    )