mohdbelal010/SecureAI-Gaurd

1"""2inference.py — SecureAI-Guard baseline inference script.3 4Reads environment variables:5    API_BASE_URL  : base URL of the SecureAI-Guard environment server6    MODEL_NAME    : OpenAI-compatible model name for the LLM agent7    HF_TOKEN      : HuggingFace token (optional, passed to model calls)8 9Logging format (required by OpenEnv):10    [START] ...episode metadata...11    [STEP]  ...per-step data...12    [END]   ...episode summary...13 14Usage:15    export API_BASE_URL=http://localhost:786016    export MODEL_NAME=gpt-3.5-turbo17    export HF_TOKEN=hf_...18    python inference.py19"""20 21import json22import logging23import os24import sys25import time26from typing import Any, Dict, Optional27 28import requests29 30# ---------------------------------------------------------------------------31# Logging setup — plain stdout so automated validators can parse the tags32# ---------------------------------------------------------------------------33logging.basicConfig(34    level=logging.INFO,35    format="%(message)s",36    handlers=[logging.StreamHandler(sys.stdout)],37)38logger = logging.getLogger("inference")39 40# ---------------------------------------------------------------------------41# Configuration from environment variables42# ---------------------------------------------------------------------------43API_BASE_URL: str = os.environ.get("API_BASE_URL", "http://localhost:7860")44MODEL_NAME: str = os.environ.get("MODEL_NAME", "gpt-3.5-turbo")45HF_TOKEN: str = os.environ.get("HF_TOKEN", "")46 47TASKS = ["basic_security", "trust_management", "adversarial_drift"]48EPISODES_PER_TASK = int(os.environ.get("EPISODES_PER_TASK", "1"))49SEED_BASE = int(os.environ.get("SEED_BASE", "42"))50 51 52# ---------------------------------------------------------------------------53# OpenAI-compatible LLM client54# ---------------------------------------------------------------------------55def call_llm(prompt: str, system: str = "") -> str:56    """57    Call an OpenAI-compatible endpoint.58    Falls back to a deterministic rule-based decision when the endpoint59    is unavailable (so the script is always runnable end-to-end).60    """61    openai_base = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")62    openai_key = os.environ.get("OPENAI_API_KEY", "")63 64    if openai_key:65        try:66            from openai import OpenAI67 68            client = OpenAI(api_key=openai_key, base_url=openai_base)69            messages = []70            if system:71                messages.append({"role": "system", "content": system})72            messages.append({"role": "user", "content": prompt})73 74            response = client.chat.completions.create(75                model=MODEL_NAME,76                messages=messages,77                max_tokens=256,78                temperature=0.0,79            )80            return response.choices[0].message.content.strip()81        except Exception as exc:82            logger.warning("LLM call failed (%s). Using rule-based fallback.", exc)83 84    # Rule-based fallback — deterministic, no network needed85    return _rule_based_decision(prompt)86 87 88def _rule_based_decision(prompt: str) -> str:89    """Deterministic fallback agent using risk keywords."""90    p = prompt.lower()91 92    HIGH_RISK = [93        "click here", "verify your account", "suspended", "claim now",94        "download", "install", "wire transfer", "credentials", "emergency",95        "send money", "customs fee", "stranded",96    ]97    MEDIUM_RISK = [98        "http://", "https://", "link", "verify", "update", "confirm",99        "invoice", "document", "offer", "free", "loan", "prize",100    ]101 102    high = sum(1 for kw in HIGH_RISK if kw in p)103    medium = sum(1 for kw in MEDIUM_RISK if kw in p)104 105    if high >= 2 or (high >= 1 and medium >= 2):106        decision, confidence, reasoning = (107            "block",108            0.92,109            (110                f"Detected {high} high-risk and {medium} medium-risk indicators. "111                "Blocking to protect user from phishing/social-engineering attempt."112            ),113        )114    elif high == 1 or medium >= 3:115        decision, confidence, reasoning = (116            "warn",117            0.75,118            (119                f"Detected {high} high-risk and {medium} medium-risk indicators. "120                "Issuing warning while preserving user trust."121            ),122        )123    elif medium >= 1:124        decision, confidence, reasoning = (125            "investigate",126            0.60,127            (128                f"Detected {medium} medium-risk indicators. "129                "Flagging for investigation — insufficient evidence to block."130            ),131        )132    else:133        decision, confidence, reasoning = (134            "allow",135            0.88,136            "No risk indicators detected. Message appears legitimate and safe to allow.",137        )138 139    return json.dumps(140        {"decision": decision, "confidence": confidence, "reasoning": reasoning}141    )142 143 144# ---------------------------------------------------------------------------145# Prompt builder146# ---------------------------------------------------------------------------147SYSTEM_PROMPT = (148    "You are an autonomous security assistant protecting a user's communications. "149    "Analyse the incoming message and decide how to handle it.\n\n"150    "Respond ONLY with valid JSON (no markdown fences):\n"151    '{"decision": "allow|block|warn|investigate", "confidence": 0.0-1.0, "reasoning": "..."}'152)153 154 155def build_prompt(obs: Dict[str, Any]) -> str:156    return (157        f"Channel     : {obs['channel']}\n"158        f"Sender      : {obs['sender']}\n"159        f"Content     : {obs['content']}\n"160        f"Risk Score  : {obs['hf_risk_score']:.2f}\n"161        f"User Trust  : {obs['user_trust']:.1f}/100\n"162        f"Sys Fatigue : {obs['system_fatigue']:.1f}/100\n"163        f"Threat Hist : {len(obs.get('threat_history', []))} recent events\n\n"164        "What is your security decision?"165    )166 167 168# ---------------------------------------------------------------------------169# Environment helpers170# ---------------------------------------------------------------------------171def env_reset(task_id: str, seed: int) -> Dict[str, Any]:172    url = f"{API_BASE_URL}/reset"173    resp = requests.post(url, json={"task_id": task_id, "seed": seed}, timeout=30)174    resp.raise_for_status()175    return resp.json()176 177 178def env_step(action: Dict[str, Any]) -> Dict[str, Any]:179    url = f"{API_BASE_URL}/step"180    resp = requests.post(url, json={"action": action}, timeout=30)181    resp.raise_for_status()182    return resp.json()183 184 185def parse_action(llm_output: str) -> Dict[str, Any]:186    """Parse LLM JSON output into an action dict."""187    # Strip markdown fences if present188    cleaned = llm_output.strip().strip("```json").strip("```").strip()189    try:190        data = json.loads(cleaned)191    except json.JSONDecodeError:192        # Last resort: default to investigate193        data = {194            "decision": "investigate",195            "confidence": 0.5,196            "reasoning": "Unable to parse LLM output; defaulting to investigation.",197        }198 199    # Validate / clamp200    valid_decisions = {"allow", "block", "warn", "investigate"}201    if data.get("decision") not in valid_decisions:202        data["decision"] = "investigate"203    data["confidence"] = float(max(0.0, min(1.0, data.get("confidence", 0.5))))204    if not data.get("reasoning", "").strip():205        data["reasoning"] = "No reasoning provided."206    return data207 208 209# ---------------------------------------------------------------------------210# Main inference loop211# ---------------------------------------------------------------------------212def run_episode(task_id: str, seed: int, episode_num: int) -> Dict[str, Any]:213    reset_data = env_reset(task_id, seed)214    obs = reset_data["observation"]215 216    episode_summary: Dict[str, Any] = {217        "task_id": task_id,218        "seed": seed,219        "episode": episode_num,220        "steps": [],221        "total_reward": 0.0,222        "final_score": None,223        "grade": None,224    }225 226    logger.info(227        "[START] task=%s episode=%d seed=%d model=%s api=%s",228        task_id,229        episode_num,230        seed,231        MODEL_NAME,232        API_BASE_URL,233    )234 235    step_num = 0236    done = False237 238    while not done:239        step_num += 1240 241        # Build prompt and get action from LLM242        prompt = build_prompt(obs)243        llm_output = call_llm(prompt, system=SYSTEM_PROMPT)244        action = parse_action(llm_output)245 246        # Step environment247        step_data = env_step(action)248        reward_val = step_data["reward"]["value"]249        episode_summary["total_reward"] += reward_val250        done = step_data["done"]251 252        step_log = {253            "step": step_num,254            "channel": obs["channel"],255            "sender": obs["sender"],256            "decision": action["decision"],257            "confidence": action["confidence"],258            "reward": reward_val,259            "user_trust": step_data["state"]["user_trust"],260            "system_fatigue": step_data["state"]["system_fatigue"],261            "threat_type": step_data["info"].get("threat_type", "unknown"),262            "done": done,263        }264        episode_summary["steps"].append(step_log)265 266        logger.info(267            "[STEP] step=%d decision=%s confidence=%.2f reward=%.4f "268            "trust=%.1f fatigue=%.1f threat=%s",269            step_num,270            action["decision"],271            action["confidence"],272            reward_val,273            step_data["state"]["user_trust"],274            step_data["state"]["system_fatigue"],275            step_data["info"].get("threat_type", "unknown"),276        )277 278        # Advance observation279        obs = step_data["observation"]280 281        # Retrieve grade if episode ended282        if done and "grade" in step_data:283            grade_data = step_data["grade"]284            episode_summary["final_score"] = grade_data.get("score")285            episode_summary["grade"] = grade_data.get("grade")286 287    episode_summary["total_reward"] = round(episode_summary["total_reward"], 4)288 289    logger.info(290        "[END] task=%s episode=%d steps=%d total_reward=%.4f score=%s grade=%s",291        task_id,292        episode_num,293        step_num,294        episode_summary["total_reward"],295        episode_summary.get("final_score"),296        episode_summary.get("grade"),297    )298 299    return episode_summary300 301 302def main():303    logger.info("=== SecureAI-Guard Inference ===")304    logger.info("API_BASE_URL : %s", API_BASE_URL)305    logger.info("MODEL_NAME   : %s", MODEL_NAME)306    logger.info("HF_TOKEN     : %s", "set" if HF_TOKEN else "not set")307 308    all_results = []309    global_episode = 0310 311    for task_id in TASKS:312        for ep_idx in range(EPISODES_PER_TASK):313            global_episode += 1314            seed = SEED_BASE + global_episode315            try:316                summary = run_episode(task_id, seed, global_episode)317                all_results.append(summary)318            except Exception as exc:319                logger.error("Episode failed: task=%s episode=%d error=%s", task_id, global_episode, exc)320 321    # Aggregate summary322    if all_results:323        avg_reward = sum(r["total_reward"] for r in all_results) / len(all_results)324        scored = [r for r in all_results if r["final_score"] is not None]325        avg_score = sum(r["final_score"] for r in scored) / len(scored) if scored else None326        logger.info(327            "=== SUMMARY === episodes=%d avg_reward=%.4f avg_score=%s",328            len(all_results),329            avg_reward,330            f"{avg_score:.4f}" if avg_score is not None else "n/a",331        )332 333 334if __name__ == "__main__":335    main()336