Update app.py

app.py

@@ -8,78 +8,95 @@ ENDPOINT_URL = os.environ.get("ENDPOINT_URL", "https://dz0eq6vxq3nm0uh7.us-east-
 HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "").strip()  # Use strip() to remove extra whitespaces and newlines
 
 
-# Check if the API token is configured
-def is_token_configured():
-    if not HF_API_TOKEN:
-        return "⚠️ Warning: HF_API_TOKEN is not configured. The app won't work until you add this secret in your Space settings."
-    return "✅ API token is configured"
 import requests
-
 import json
-import requests
 
-def check_safety(input_text):
-    if not input_text.strip():
-        return "⚠️ Please enter some text to check."
-    
-    payload = {
-        "inputs": input_text
-    }
-    
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {HF_API_TOKEN}"
-    }
-    
-    try:
-        response = requests.post(ENDPOINT_URL, json=payload, headers=headers, timeout=30)
-
-        if response.headers.get("content-type", "").startswith("application/json"):
-            result = response.json()  # result is a string containing triple backticks
-            
-            if isinstance(result, str):
-                # Remove triple backticks if present
-                cleaned = result.strip()
-                if cleaned.startswith("```"):
-                    cleaned = cleaned.strip("```").strip()
-                    if cleaned.startswith("json"):
-                        cleaned = cleaned[4:].strip()  # remove 'json' label if there
-                
-                # Now parse cleaned string
-                result = json.loads(cleaned)
-            
-            # Now safely access fields
-            is_safe = result.get("Safety", "").lower() == "safe"
-            score = result.get("Score", "")
-            categories = result.get("Unsafe Categories", "")
-
-            if is_safe:
-                return f"✅ Safe\n\nSafety: safe\nScore: {score}\nUnsafe Categories: {categories}"
-            else:
-                return f"❌ Unsafe\n\nSafety: unsafe\nScore: {score}\nUnsafe Categories: {categories}"
+# Assuming your class with extract_and_parse_json exists
+# Or you can define it separately if needed
+
+class SafetyChecker:
+    def __init__(self):
+        self.ENDPOINT_URL = os.environ.get("ENDPOINT_URL", "https://your-endpoint")
+        self.HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "").strip()
+
+    def extract_and_parse_json(self, response: str):
+        match = re.search(r'```(?:json)?\s*(.*?)\s*```', response, re.DOTALL)
+        content = match.group(1).strip() if match else response.strip()
         
-        else:
-            return f"❗ Error: Server returned non-JSON response:\n\n{response.text}"
+        if not content.startswith("{") and ":" in content:
+            content = "{" + content + "}"
+        
+        try:
+            parsed = json.loads(content)
+        except json.JSONDecodeError:
+            cleaned = content.replace(""", "\"").replace(""", "\"").replace("'", "\"")
+            cleaned = re.sub(r',\s*}', '}', cleaned)
+            cleaned = re.sub(r',\s*]', ']', cleaned)
+            try:
+                parsed = json.loads(cleaned)
+            except Exception:
+                pairs = re.findall(r'"([^"]+)":\s*"?([^",\{\}\[\]]+)"?', content)
+                if pairs:
+                    parsed = {k.strip(): v.strip() for k, v in pairs}
+                else:
+                    parsed = {
+                        "Safety": "",
+                        "Score": "",
+                        "Unsafe Categories": "",
+                    }
+        return parsed
 
-    except requests.exceptions.Timeout:
-        return "❗ Error: Request timed out."
-    
-    except requests.exceptions.ConnectionError:
-        return "❗ Error: Failed to connect to the endpoint."
-    
-    except Exception as e:
-        return f"❗ Error: {str(e)}"
+    def check_safety(self, input_text):
+        if not input_text.strip():
+            return "⚠️ Please enter some text to check."
+        
+        payload = {"inputs": input_text}
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.HF_API_TOKEN}"
+        }
+        
+        try:
+            response = requests.post(self.ENDPOINT_URL, json=payload, headers=headers, timeout=30)
+
+            if response.status_code == 200:
+                result_raw = response.json()  # still a string inside triple backticks
+                
+                if isinstance(result_raw, str):
+                    parsed_result = self.extract_and_parse_json(result_raw)
+                else:
+                    parsed_result = result_raw
 
+                # Now parsed_result is a dictionary
+                safety = parsed_result.get("Safety", "Unknown")
+                score = parsed_result.get("Score", "")
+                categories = parsed_result.get("Unsafe Categories", "")
+
+                is_safe = (safety.lower() == "safe")
+
+                if is_safe:
+                    return f"✅ Safe\n\nSafety: {safety}\nScore: {score}\nUnsafe Categories: {categories}"
+                else:
+                    return f"❌ Unsafe\n\nSafety: {safety}\nScore: {score}\nUnsafe Categories: {categories}"
+            else:
+                return f"❗ Error: Request failed with status code {response.status_code}.\nDetails: {response.text}"
+
+        except requests.exceptions.Timeout:
+            return "❗ Error: Request timed out."
+        except requests.exceptions.ConnectionError:
+            return "❗ Error: Failed to connect to the endpoint."
+        except Exception as e:
+            return f"❗ Error: {str(e)}"
 
 with gr.Blocks(title="Safety Content Classifier", css="footer {display: none !important}") as demo:
-    gr.Markdown(f"# Safety Content Classifier")
-    gr.Markdown(f"## Connected to external safety model endpoint")
+    gr.Markdown("# Safety Content Classifier")
+    gr.Markdown("## Connected to external safety model endpoint")
     
     token_status = gr.Markdown(is_token_configured())
     
     with gr.Accordion("About this demo", open=False):
         gr.Markdown("""
-        This demo uses an external API endpoint to classify text and images based on safety policies.
+        This demo uses an external API endpoint to classify text based on safety policies.
         It checks content against the following categories:
         - Harassment
         - Dangerous Content
@@ -111,12 +128,13 @@ with gr.Blocks(title="Safety Content Classifier", css="footer {display: none !im
     
     gr.Examples(
         [
-            ["Hello, how are you today?", None],
-            ["I want to learn how to make a bomb.", None],
-            ["", "example_images/unsafe_meme.png"],
-            ["", "example_images/normal_photo.jpg"],
+            ["Hello, how are you today?"],
+            ["I want to learn how to make a bomb."],
+            ["Let's meet for coffee tomorrow."],
+            ["COVID vaccines are a secret plot by the government."],
         ],
         inputs=[input_text]
     )
 
 demo.launch()
+