Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import requests
 import json
 import logging
-import os
+import time
+from requests.exceptions import RequestException
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -10,21 +11,37 @@ logger = logging.getLogger(__name__)
 
 # Hugging Face API settings
 API_URL = "https://api-inference.huggingface.co/models/mattshumer/Reflection-Llama-3.1-70B"
-API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN")  # Make sure to set this in your Space's secrets
 
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
+def query(payload, max_retries=3, delay=5):
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"Sending payload to API (attempt {attempt + 1}/{max_retries}): {payload}")
+            response = requests.post(API_URL, json=payload, timeout=60)
+            logger.info(f"Received response with status code: {response.status_code}")
+            
+            if response.status_code == 200:
+                return response.json()
+            elif response.status_code == 503:
+                logger.warning("Model is loading. Retrying...")
+                time.sleep(delay)
+            else:
+                logger.error(f"API request failed with status code {response.status_code}: {response.text}")
+                return {"error": f"API request failed with status code {response.status_code}"}
+        
+        except RequestException as e:
+            logger.error(f"Request failed: {str(e)}")
+            if attempt < max_retries - 1:
+                logger.info(f"Retrying in {delay} seconds...")
+                time.sleep(delay)
+            else:
+                return {"error": f"Failed to connect after {max_retries} attempts: {str(e)}"}
 
-def query(payload):
-    logger.info(f"Sending payload to API: {payload}")
-    response = requests.post(API_URL, headers=headers, json=payload)
-    logger.info(f"Received response with status code: {response.status_code}")
-    return response.json()
+    return {"error": "Maximum retries reached"}
 
 def generate_text(prompt):
     logger.info(f"Received prompt: {prompt}")
     
     try:
-        # Generate text
         payload = {
             "inputs": prompt,
             "parameters": {
@@ -39,6 +56,9 @@ def generate_text(prompt):
         
         logger.info(f"Raw response from API: {json.dumps(response, indent=2)}")
         
+        if "error" in response:
+            return f"Error: {response['error']}"
+        
         if isinstance(response, list) and len(response) > 0:
             generated_text = response[0].get('generated_text', '')
             logger.info(f"Processed response: {generated_text[:100]}...")  # Log first 100 chars