""" inference.py — agentic_security_lab/inference.py (project root) ================================================================= Mandatory stdout format: [START] task= env=agentic-security-lab model= [STEP] step= action= reward=<0.00> done= error= [END] success= steps= score=<0.00> rewards= Env vars (set in HF Space secrets): API_BASE_URL LLM endpoint — default: HF inference router MODEL_NAME model id — default: Qwen/Qwen2.5-7B-Instruct HF_TOKEN API key — injected automatically by the validator ENV_BASE_URL Space URL — default: localhost (validator sets this) TASK_NAME easy|medium|hard (omit to run all three) """ import json import os import textwrap from typing import Optional import httpx from openai import OpenAI # ── Configuration ────────────────────────────────────────────────────────────── # Defaults match what the hackathon validator injects. # HF_TOKEN is set automatically by the validator — do not hardcode it. API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct") API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or "" ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:8000") ALL_TASKS = ["easy", "medium", "hard"] MAX_STEPS = 30 BENCHMARK = "agentic-security-lab" SYSTEM_PROMPT = textwrap.dedent(""" You are an expert security engineer responding to a supply-chain incident. Issue ONE action per turn as a JSON object with keys "command" and "parameters". Commands: inspect_package {"package": ""} check_dependents {"package": ""} rotate_secret {"secret": ""} quarantine {"package": ""} notify {"team": ""} scan_logs {"package": ""} conclude {} Strategy: 1. scan_logs / inspect_package on all packages to identify IOCs. 2. quarantine every malicious package immediately. 3. check_dependents to find all affected downstream teams. 4. rotate_secret for every exposed credential — critical ones first. 5. notify every affected team. 6. conclude when done. Reply with ONLY the JSON object. No markdown, no explanation, no prose. Example: {"command": "quarantine", "parameters": {"package": "axios@1.7.4"}} """).strip() # ── Logging helpers ──────────────────────────────────────────────────────────── def log_start(task: str, model: str) -> None: print(f"[START] task={task} env={BENCHMARK} model={model}", flush=True) def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None: print( f"[STEP] step={step} action={action} " f"reward={reward:.2f} done={str(done).lower()} " f"error={error if error else 'null'}", flush=True, ) def log_end(success: bool, steps: int, score: float, rewards: list) -> None: rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={str(success).lower()} steps={steps} " f"score={score:.2f} rewards={rewards_str}", flush=True, ) # ── Environment helpers ──────────────────────────────────────────────────────── def env_reset(http: httpx.Client, task_name: str) -> dict: r = http.post("/reset", json={"task_name": task_name}) r.raise_for_status() return r.json() def env_step(http: httpx.Client, command: str, parameters: dict) -> dict: r = http.post("/step", json={"command": command, "parameters": parameters}) r.raise_for_status() return r.json() # ── Action parsing ───────────────────────────────────────────────────────────── def parse_action(raw: str) -> tuple: text = raw.strip() if text.startswith("```"): lines = text.splitlines() text = "\n".join(l for l in lines if not l.startswith("```")).strip() try: obj = json.loads(text) return obj.get("command", "conclude"), obj.get("parameters", {}) except (json.JSONDecodeError, AttributeError): return "conclude", {} # ── Episode runner ───────────────────────────────────────────────────────────── def run_task(llm: OpenAI, http: httpx.Client, task_name: str) -> float: rewards = [] step_n = 0 success = False score = 0.0 log_start(task=task_name, model=MODEL_NAME) try: obs = env_reset(http, task_name) messages = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": obs["result"]}, ] for step_n in range(1, MAX_STEPS + 1): # ── LLM call — wrapped so a transient API error doesn't crash ── try: resp = llm.chat.completions.create( model = MODEL_NAME, messages = messages, max_tokens = 150, temperature = 0.2, ) raw = resp.choices[0].message.content or "" except Exception as llm_err: # Treat LLM failure as a no-op conclude to end gracefully raw = '{"command": "conclude", "parameters": {}}' print(f"[DEBUG] LLM error at step {step_n}: {llm_err}", flush=True) command, params = parse_action(raw) action_str = json.dumps({"command": command, "parameters": params}) # ── Environment step ─────────────────────────────────────────── obs = env_step(http, command, params) reward = float(obs.get("reward", 0.0)) done = bool(obs.get("done", False)) error = obs.get("error") or None rewards.append(reward) log_step(step=step_n, action=action_str, reward=reward, done=done, error=error) messages.append({"role": "assistant", "content": raw}) messages.append({"role": "user", "content": obs["result"]}) if done: break score = min(1.0, max(0.0, sum(rewards))) success = score > 0.0 except Exception as exc: # Always emit [END] — even on unexpected failure log_end(success=False, steps=step_n, score=0.0, rewards=rewards) raise log_end(success=success, steps=step_n, score=score, rewards=rewards) return score # ── Entry point ──────────────────────────────────────────────────────────────── def main() -> None: llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) http = httpx.Client(base_url=ENV_BASE_URL, timeout=60) task_env = os.getenv("TASK_NAME", "") tasks = [task_env] if task_env in ALL_TASKS else ALL_TASKS scores = {} try: for task in tasks: scores[task] = run_task(llm, http, task) finally: http.close() print("\n=== Baseline Results ===", flush=True) for task, s in scores.items(): print(f" {task:8s} score={s:.2f}", flush=True) if __name__ == "__main__": main()