Testing the complaince engine fix

Author: PrashantMaht0

ai-service/app.py

@@ -1,9 +1,9 @@
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 import requests
 import os
+import json
 from dotenv import load_dotenv
 from compliance_engine import ComplianceAuditor
-import json
 
 # Load your Hugging Face token from the .env file in ai-service/
 load_dotenv()
@@ -21,28 +21,63 @@ async def run_audit(
     boxes: str = Form(...)
 ):
     try:
-        # 1. Package the incoming data to forward to Hugging Face
         headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+        
+        # Read the image into memory ONCE so we can reuse it for multiple chunks
         image_bytes = await screenshot.read()
-        files = {"screenshot": (screenshot.filename, image_bytes, screenshot.content_type)}
-        data = {"words": words, "boxes": boxes}
+        
+        # Parse the incoming JSON strings into Python lists
+        words_list = json.loads(words)
+        boxes_list = json.loads(boxes)
+        
+        # --- THE SLIDING WINDOW CHUNKING LOGIC ---
+        CHUNK_SIZE = 400
+        all_flagged_predictions = []
+        
+        total_elements = len(words_list)
+        print(f"[*] Starting Sliding Window analysis for {total_elements} elements on {target_url}...")
+
+        # Loop through the page in batches of 400 to prevent HF from hitting the 512 token limit
+        for i in range(0, total_elements, CHUNK_SIZE):
+            chunk_words = words_list[i : i + CHUNK_SIZE]
+            chunk_boxes = boxes_list[i : i + CHUNK_SIZE]
+            
+            if not chunk_words:
+                continue
+                
+            print(f"[*] Sending Chunk to Hugging Face: {i} to {i + len(chunk_words)}...")
+            
+            # Package this specific chunk to send to Hugging Face
+            files = {"screenshot": (screenshot.filename, image_bytes, screenshot.content_type)}
+            data = {
+                "words": json.dumps(chunk_words), 
+                "boxes": json.dumps(chunk_boxes)
+            }
 
-        print(f"[*] Analyzing {target_url} via Hugging Face Cloud...")
-        hf_response = requests.post(HF_API_URL, headers=headers, files=files, data=data)
+            hf_response = requests.post(HF_API_URL, headers=headers, files=files, data=data)
 
-        if hf_response.status_code != 200:
-            raise Exception(f"Hugging Face API Error: {hf_response.text}")
+            if hf_response.status_code != 200:
+                print(f"[!] Warning: HF API Error on chunk {i}: {hf_response.text}")
+                continue # Skip this chunk if HF throws a timeout or error, but keep processing the rest of the page!
 
-        ai_predictions = hf_response.json()
+            chunk_predictions = hf_response.json()
+            
+            # Combine the flagged items from this chunk into our master list
+            if isinstance(chunk_predictions, list):
+                all_flagged_predictions.extend(chunk_predictions)
 
+        print(f"[+] Chunking complete! Found {len(all_flagged_predictions)} total suspicious elements to audit.")
 
-        # 2. Pass the AI predictions into your Compliance Engine
-        print("[*] Generating Legal Compliance Report...")
+        # --- STAGE 2: THE GEMINI LOGICAL AUDIT ---
+        print("[*] Generating Legal Compliance Report via Gemini...")
         auditor = ComplianceAuditor(target_url=target_url)
-        final_report = auditor.analyze_detections(ai_predictions)
+        
+        # We pass the MASSIVE stitched list of all found patterns to your Gemini engine
+        final_report = auditor.analyze_detections(all_flagged_predictions)
 
         # 3. Return the formatted JSON to your Node.js backend
         return final_report
 
     except Exception as e:
+        print(f"[!] Fatal Audit Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))

ai-service/compliance_engine.py

@@ -7,9 +7,12 @@
 
 load_dotenv()
 
-# Initialize the stable google-generativeai SDK
+# Initialize the SDK and force strict JSON output to prevent parsing crashes
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
-llm_model = genai.GenerativeModel('gemini-2.5-flash')
+llm_model = genai.GenerativeModel(
+    model_name='gemini-2.5-flash',
+    generation_config={"response_mime_type": "application/json"}
+)
 
 REGULATORY_MAP = {
     "preselected_invasive_default": {
@@ -81,46 +84,43 @@ def _verify_with_llm(self, element_text, layout_label):
         """The Stage 2 Classifier: Asks Gemini to legally categorize the text using dynamic map data."""
 
         legal_framework = "\n".join([f"{i+1}. {k.upper()}: {v['description']}" for i, (k, v) in enumerate(REGULATORY_MAP.items())])
-        
         allowed_categories = "[" + ", ".join(REGULATORY_MAP.keys()) + ", safe]"
 
         prompt = f"""
-        ACT AS: A Senior Digital Rights Attorney and GDPR Auditor specialized in the Digital Services Act (DSA) Article 25 & 27.
+        ACT AS: A Senior Digital Rights Attorney and GDPR Auditor.
         
-        TASK: Conduct a high-stakes audit on a specific UI element to determine if it constitutes a "Dark Pattern" (deceptive design).
+        TASK: Evaluate this specific UI element text flagged by a Vision AI. Is it a genuine "Dark Pattern" (deceptive design) or a normal UI element?
         
         CONTEXT:
         - Element Type: {layout_label}
         - Detected Text: "{element_text}"
         
-        LEGAL REFERENCE FRAMEWORK:
+        LEGAL FRAMEWORK (CATEGORIES):
         {legal_framework}
 
+        CRITICAL: NEGATIVE EXAMPLES (IGNORE THESE - LABEL AS 'safe')
+        - Standard navigation ("Home", "About Us", "Contact").
+        - Standard actions ("Login", "Submit", "Search", "Read More", "Accept").
+        - Cookie banners with fair choices ("Accept All" alongside "Decline All").
+        - "No thanks" or "Close" buttons.
+
         AUDIT RULES:
-        - ZERO TOLERANCE FOR FALSE POSITIVES: If the text is standard, polite, or merely descriptive (e.g., "We use cookies", "Learn More", "Accept"), it MUST be labeled 'safe'.
+        - ZERO TOLERANCE FOR FALSE POSITIVES: If the text is standard, polite, or merely descriptive, it MUST be labeled 'safe'.
         - CONTEXT MATTERS: "No thanks" is safe. "No, I prefer to pay more" is emotional_steering.
-        - DEFAULT TO SAFE: If you are less than 95% certain a pattern exists, return 'safe'.
 
-        STEP-BY-STEP REASONING:
-        1. Analyze the literal meaning of the text.
-        2. Evaluate the psychological intent (Is it steering, shaming, or confusing?).
-        3. Compare against the legal frameworks above.
-        
         OUTPUT FORMAT:
-        You must return a raw JSON object with this exact structure:
+        Return ONLY a JSON object. You MUST provide the "reasoning" key BEFORE the "category" key to ensure logical chain-of-thought analysis.
         {{
-            "reasoning": "A 1-sentence legal justification for your decision.",
+            "reasoning": "CHAIN OF THOUGHT: Step-by-step, logically explain why this text violates user intent OR why it is perfectly safe.",
             "category": "one_of_the_categories_below_or_safe"
         }}
 
-        CATEGORIES:
-        {allowed_categories}
+        ALLOWED CATEGORIES: {allowed_categories}
         """
 
         try:
             response = llm_model.generate_content(prompt)
-            result_text = response.text.replace('```json', '').replace('```', '').strip()
-            data = json.loads(result_text)
+            data = json.loads(response.text) # Clean parsing since response_mime_type is JSON
 
             category = data.get("category", "safe")
             reasoning = data.get("reasoning", "No explanation provided.")
@@ -133,21 +133,29 @@ def _verify_with_llm(self, element_text, layout_label):
 
     def analyze_detections(self, hf_api_response):
         """Orchestrates the two-stage pipeline."""
-        ai_predictions = hf_api_response.get("flagged_elements", [])
-        print("\n[*] Running Stage 2 LLM Classification on flagged elements...")
+        
+        # Accommodates both old list structure and new dictionary structure
+        if isinstance(hf_api_response, list):
+            ai_predictions = hf_api_response
+        else:
+            ai_predictions = hf_api_response.get("flagged_elements", [])
+            
+        print(f"\n[*] Running Stage 2 LLM Classification on {len(ai_predictions)} flagged elements...")
 
         for detection in ai_predictions:
-            layout_label = detection.get("predicted_label") 
-            bbox = detection.get("box_2d") or [0, 0, 0, 0] 
+            # Safely extract data handling potential missing keys from Hugging Face output
+            layout_label = detection.get("layoutlm_label", "deceptive_element") 
+            bbox = detection.get("box_2d", [0, 0, 0, 0])
             element_text = detection.get("text", "") 
 
-            if layout_label in ["action_button", "overlay_content", "deceptive_element"]:
+            if element_text:
                 print(f"[*] Auditing text: '{element_text}'")
 
                 category, reasoning = self._verify_with_llm(element_text, layout_label)
-                time.sleep(1.5)
+                time.sleep(1.5) # Prevents Gemini API rate limiting
 
                 if category in REGULATORY_MAP:
+                    print(f"    [!] Violation Found: {category}")
                     rule = REGULATORY_MAP[category]
                     self.trust_score -= rule["penalty"]