hch-dev's picture
Update Version_5/app.py
0ea26e0 verified
Raw
History Blame Contribute Delete
4.33 kB
import time
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field
# Absolute path imports mapped to your Mega-Server setup
from Version_5.src.dom_scraper import extract_dom_features
from Version_5.src.sub_agents import (
agent_url_analyst,
agent_html_structure,
agent_content_semantics,
agent_brand_impersonation
)
from Version_5.src.orchestrator import evaluate_consensus, run_judge
router = APIRouter()
# Define API request payload structure
class ThreatAnalysisRequest(BaseModel):
url: str = Field(..., description="Target landing page URL to analyze", example="http://example.com")
sender: str = Field(..., description="Alleged sender address header", example="security@paypal.com")
email_body: str = Field(..., description="Full text/body payload of the incoming message")
@router.post("/predict")
async def analyze_payload_endpoint(payload: ThreatAnalysisRequest):
url = payload.url.strip()
sender = payload.sender.strip()
email_body = payload.email_body.strip()
if not url and not email_body:
raise HTTPException(
status_code=400,
detail="Provide at least a validation URL or a message body."
)
try:
start_time = time.perf_counter()
# 1. Graceful Bypass for Missing URLs
if url == "":
url_report = "Safe/Neutral: No URL was present in the email to evaluate."
html_report = "Safe/Neutral: No web page to evaluate."
else:
dom_data = extract_dom_features(url)
url_report = agent_url_analyst(url)
html_report = agent_html_structure(dom_data)
# 2. Graceful Bypass for Missing Sender
if sender == "":
brand_report = "Neutral: No sender address provided to verify brand impersonation."
else:
brand_report = agent_brand_impersonation(email_body, sender)
# 3. Multi-Specialist Forensic Panel
reports = {
"URL_Agent": url_report,
"HTML_Agent": html_report,
"Content_Agent": agent_content_semantics(email_body),
"Brand_Agent": brand_report
}
# ... (Consensus and Judge logic remains the same below)
# Tier 3: Core Consensus Evaluator
consensus_victory = evaluate_consensus(reports)
if consensus_victory:
# Handle standard serialization if agents return Pydantic objects or plain strings
final_verdict = {
"verdict": reports["URL_Agent"].claim if hasattr(reports["URL_Agent"], "claim") else str(reports["URL_Agent"]),
"confidence_score": reports["URL_Agent"].confidence if hasattr(reports["URL_Agent"], "confidence") else 1.0,
"justification": "Bypassed judicial review due to absolute sub-agent unanimity across forensics."
}
else:
# Tier 4: Judicial Override via Groq Cloud API
reports_str = "\n".join([
f"[{name}]\n{r.model_dump_json(indent=2) if hasattr(r, 'model_dump_json') else str(r)}"
for name, r in reports.items()
])
raw_data = f"Target URL: {url}\nTarget Sender: {sender}\nBody: {email_body}"
judge_verdict = run_judge(reports_str, raw_data)
final_verdict = judge_verdict.model_dump() if hasattr(judge_verdict, "model_dump") else judge_verdict
latency_ms = (time.perf_counter() - start_time) * 1000
# Safe response serialization
serializable_reports = {}
for name, report in reports.items():
serializable_reports[name] = report.model_dump() if hasattr(report, "model_dump") else str(report)
return {
"target_url": url,
"target_sender": sender,
"consensus_reached": consensus_victory,
"latency_ms": round(latency_ms, 2),
"sub_agent_claims": serializable_reports,
"final_evaluation": final_verdict
}
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"Internal agent execution lifecycle crash: {str(e)}"
)