Update Version_5/src/sub_agents.py

Version_5/src/sub_agents.py

@@ -1,52 +1,61 @@
-import os
-from groq import Groq
-from .schemas import AgentReport
-
-# Initialize the Groq Cloud Client
-client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
-
-def run_agent(role_prompt: str, target_data: str) -> AgentReport:
-    # We call Groq's hosted 8B infrastructure instead of local resources
-    response = client.chat.completions.create(
-        model="llama-3.1-8b-instant",
-        messages=[
-            {"role": "system", "content": role_prompt},
-            {"role": "user", "content": target_data}
-        ],
-        # Groq enforces JSON mode via this parameter block
-        response_format={"type": "json_object"},
-        temperature=0.0
-    )
-    return AgentReport.model_validate_json(response.choices[0].message.content)
-
-def agent_url_analyst(url: str) -> AgentReport:
-    prompt = (
-        "You are a URL Shadows Analyst. Respond strictly in valid JSON matching this schema:\n"
-        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
-        "Analyze the structure for structural entropy, unusual domains, and brand keywords."
-    )
-    return run_agent(prompt, f"URL to analyze: {url}")
-
-def agent_html_structure(dom_json: str) -> AgentReport:
-    prompt = (
-        "You are an HTML Code Analyst. Respond strictly in valid JSON matching this schema:\n"
-        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
-        "Analyze these DOM elements. Flag external tracking inputs or forms posting to alternative servers."
-    )
-    return run_agent(prompt, f"DOM Data: {dom_json}")
-
-def agent_content_semantics(email_body: str) -> AgentReport:
-    prompt = (
-        "You are a Phishing Copywriter Critic. Respond strictly in valid JSON matching this schema:\n"
-        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
-        "Extract semantic compliance anomalies, high emotional coercion markers, and urgency loops."
-    )
-    return run_agent(prompt, f"Email Body: {email_body}")
-
-def agent_brand_impersonation(email_body: str, sender: str) -> AgentReport:
-    prompt = (
-        "You are an Identity Protection Agent. Respond strictly in valid JSON matching this schema:\n"
-        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
-        "Flag mismatched target namespaces where sender domains do not match corporate identifiers."
-    )
+import os
+from groq import Groq
+from .schemas import AgentReport
+
+# Initialize the Groq Cloud Client
+client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
+# ── NEW: System Override for AMTSO & Simulated Tests ──
+SIMULATION_OVERRIDE = """
+CRITICAL RULE - SIMULATED THREATS: 
+If the text, HTML, or content explicitly claims to be a "Test Page", "Dummy Page", "AMTSO Phishing Check", or states it is designed to test anti-malware/antivirus software, YOU MUST treat it as a real threat. Return a claim of "PHISHING" with a high confidence score, and explicitly state in the evidence array that it is a simulated testing environment. Do not mark it as LEGITIMATE just because it is a benign test.
+"""
+
+def run_agent(role_prompt: str, target_data: str) -> AgentReport:
+    # We call Groq's hosted 8B infrastructure instead of local resources
+    response = client.chat.completions.create(
+        model="llama-3.1-8b-instant",
+        messages=[
+            {"role": "system", "content": role_prompt},
+            {"role": "user", "content": target_data}
+        ],
+        # Groq enforces JSON mode via this parameter block
+        response_format={"type": "json_object"},
+        temperature=0.0
+    )
+    return AgentReport.model_validate_json(response.choices[0].message.content)
+
+def agent_url_analyst(url: str) -> AgentReport:
+    prompt = (
+        "You are a URL Shadows Analyst. Respond strictly in valid JSON matching this schema:\n"
+        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
+        "Analyze the structure for structural entropy, unusual domains, and brand keywords."
+    )
+    return run_agent(prompt, f"URL to analyze: {url}")
+
+def agent_html_structure(dom_json: str) -> AgentReport:
+    prompt = (
+        "You are an HTML Code Analyst. Respond strictly in valid JSON matching this schema:\n"
+        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
+        "Analyze these DOM elements. Flag external tracking inputs or forms posting to alternative servers.\n"
+        + SIMULATION_OVERRIDE
+    )
+    return run_agent(prompt, f"DOM Data: {dom_json}")
+
+def agent_content_semantics(email_body: str) -> AgentReport:
+    prompt = (
+        "You are a Phishing Copywriter Critic. Respond strictly in valid JSON matching this schema:\n"
+        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
+        "Extract semantic compliance anomalies, high emotional coercion markers, and urgency loops.\n"
+        + SIMULATION_OVERRIDE
+    )
+    return run_agent(prompt, f"Email Body: {email_body}")
+
+def agent_brand_impersonation(email_body: str, sender: str) -> AgentReport:
+    prompt = (
+        "You are an Identity Protection Agent. Respond strictly in valid JSON matching this schema:\n"
+        "{\"claim\": \"PHISHING\"|\"LEGITIMATE\"|\"AMBIGUOUS\", \"confidence\": float, \"evidence\": [string]}\n"
+        "Flag mismatched target namespaces where sender domains do not match corporate identifiers.\n"
+        + SIMULATION_OVERRIDE
+    )
     return run_agent(prompt, f"Sender: {sender}\nBody: {email_body}")