""" The swarm. Four specialist roles built on Gemma-4-on-Cerebras: Analyst (fan-out, per finding) -> category, CVSS, exploitability, FP-likelihood Remediator (conditional, per High/Critical) -> fix + minimal patch hint Commander (once, over the batch) -> exec rollup + prioritized actions Vision RCA (multimodal, on demand) -> root-cause + patch from a screenshot The speed thesis lives in `run_batch`: analysts fan out concurrently (bounded by a semaphore) so the wall-clock is ~one slow call, not the sum of all calls. The sequential baseline runs the identical work one-at-a-time — same model, honest delta. """ from __future__ import annotations import asyncio import time from typing import Any, AsyncIterator, Optional import httpx # ---- strict JSON schemas (constrain Gemma 4 to clean, parseable output) ---- ANALYST_SCHEMA = { "type": "object", "additionalProperties": False, "properties": { "category": {"type": "string"}, "severity_cvss": {"type": "number"}, "severity_label": {"type": "string", "enum": ["critical", "high", "medium", "low", "info"]}, "exploitability": {"type": "string", "enum": ["trivial", "plausible", "theoretical", "none"]}, "false_positive_likelihood": {"type": "number"}, "rationale": {"type": "string"}, "affected_asset": {"type": "string"}, }, "required": ["category", "severity_cvss", "severity_label", "exploitability", "false_positive_likelihood", "rationale", "affected_asset"], } REMEDIATOR_SCHEMA = { "type": "object", "additionalProperties": False, "properties": { "fix_summary": {"type": "string"}, "patch_hint": {"type": "string"}, "effort": {"type": "string", "enum": ["S", "M", "L"]}, }, "required": ["fix_summary", "patch_hint", "effort"], } COMMANDER_SCHEMA = { "type": "object", "additionalProperties": False, "properties": { "headline": {"type": "string"}, "risk_paragraph": {"type": "string"}, "top_actions": {"type": "array", "items": {"type": "string"}}, "counts": { "type": "object", "additionalProperties": False, "properties": {k: {"type": "integer"} for k in ["critical", "high", "medium", "low", "info"]}, "required": ["critical", "high", "medium", "low", "info"], }, }, "required": ["headline", "risk_paragraph", "top_actions", "counts"], } VISION_SCHEMA = { "type": "object", "additionalProperties": False, "properties": { "root_cause": {"type": "string"}, "evidence_from_image": {"type": "string"}, "concrete_patch": {"type": "string"}, "confidence": {"type": "number"}, }, "required": ["root_cause", "evidence_from_image", "concrete_patch", "confidence"], } ANALYST_SYS = ( "You are a senior application-security analyst triaging one scanner finding. " "Judge real risk, not the scanner's raw label. Estimate a CVSS-style 0-10 score, " "whether it's actually exploitable, and how likely it's a false positive. " "Be terse and specific. Output only the requested JSON." ) REMEDIATOR_SYS = ( "You are a remediation engineer. Given a triaged finding, give the smallest correct " "fix and a one-line patch hint a developer can act on immediately. Output only JSON." ) COMMANDER_SYS = ( "You are the incident commander. Given the triaged batch, write a crisp executive " "rollup: one headline, one risk paragraph, and the top prioritized actions. " "Count findings by severity. Output only JSON." ) VISION_SYS = ( "You are a security analyst doing root-cause analysis from a screenshot of code, a " "dashboard, or an alert. Identify the actual root cause, cite what you see in the image, " "and give a concrete patch. Output only JSON." ) def _analyst_user(f: dict) -> str: return (f"scanner={f.get('scanner','?')} rule={f.get('rule','?')} " f"raw_severity={f.get('raw_severity','?')}\n" f"path={f.get('path','')}\n" f"message: {f.get('message','')}\n" f"snippet: {f.get('snippet','')[:400]}") class Swarm: def __init__(self, client, *, max_concurrency: int = 10, remediate_min_cvss: float = 7.0): self.client = client self.sem = asyncio.Semaphore(max_concurrency) self.remediate_min_cvss = remediate_min_cvss # ---- single agents ---- async def analyst(self, http: httpx.AsyncClient, f: dict) -> dict: r = await self.client.chat(http, system=ANALYST_SYS, user=_analyst_user(f), json_schema=ANALYST_SCHEMA, max_tokens=350) return {"finding": f, "analysis": r.json or {}, "timing": r.timing} async def remediator(self, http: httpx.AsyncClient, f: dict, analysis: dict) -> Optional[dict]: a = (f"finding: {f.get('message','')}\ncategory: {analysis.get('category')}\n" f"cvss: {analysis.get('severity_cvss')}\nsnippet: {f.get('snippet','')[:300]}") r = await self.client.chat(http, system=REMEDIATOR_SYS, user=a, json_schema=REMEDIATOR_SCHEMA, max_tokens=250) return r.json async def vision_rca(self, http: httpx.AsyncClient, *, context: str, image_b64: str, media_type: str) -> dict: r = await self.client.chat(http, system=VISION_SYS, user=context, image_b64=image_b64, media_type=media_type, json_schema=VISION_SCHEMA, max_tokens=400) return {"rca": r.json or {}, "timing": r.timing} async def commander(self, http: httpx.AsyncClient, triaged: list[dict]) -> dict: lines = [] for t in triaged[:60]: a = t["analysis"] lines.append(f"- [{a.get('severity_label','?')}/{a.get('severity_cvss','?')}] " f"{a.get('category','?')}: {t['finding'].get('message','')[:90]}") r = await self.client.chat(http, system=COMMANDER_SYS, user="Triaged findings:\n" + "\n".join(lines), json_schema=COMMANDER_SCHEMA, max_tokens=500) return {"rollup": r.json or {}, "timing": r.timing} # ---- one finding, full pipeline (analyst -> conditional remediator) ---- async def _one(self, http: httpx.AsyncClient, f: dict) -> dict: res = await self.analyst(http, f) a = res["analysis"] if float(a.get("severity_cvss") or 0) >= self.remediate_min_cvss: res["remediation"] = await self.remediator(http, f, a) return res # ---- batch: stream results as they complete; bounded concurrency ---- async def run_batch(self, findings: list[dict], *, parallel: bool) -> AsyncIterator[dict]: t0 = time.perf_counter() triaged: list[dict] = [] calls = 0 ttfts: list[float] = [] comp_tokens = 0 async with httpx.AsyncClient() as http: async def guarded(f): async with self.sem: return await self._one(http, f) if parallel: tasks = [asyncio.create_task(guarded(f)) for f in findings] for fut in asyncio.as_completed(tasks): res = await fut triaged.append(res) calls += 1 + (1 if "remediation" in res else 0) if res["timing"].ttft_ms: ttfts.append(res["timing"].ttft_ms) comp_tokens += res["timing"].completion_tokens yield {"type": "finding", "data": _wire(res), "elapsed_ms": (time.perf_counter() - t0) * 1000} else: for f in findings: res = await self._one(http, f) triaged.append(res) calls += 1 + (1 if "remediation" in res else 0) if res["timing"].ttft_ms: ttfts.append(res["timing"].ttft_ms) comp_tokens += res["timing"].completion_tokens yield {"type": "finding", "data": _wire(res), "elapsed_ms": (time.perf_counter() - t0) * 1000} roll = await self.commander(http, triaged) calls += 1 yield {"type": "rollup", "data": roll["rollup"], "elapsed_ms": (time.perf_counter() - t0) * 1000} wall = (time.perf_counter() - t0) * 1000 yield {"type": "summary", "data": { "wall_ms": round(wall, 1), "calls": calls, "findings": len(findings), "avg_ttft_ms": round(sum(ttfts) / len(ttfts), 1) if ttfts else None, "completion_tokens": comp_tokens, "throughput_tps": round(comp_tokens / (wall / 1000), 1) if wall else 0, "mode": "parallel" if parallel else "sequential", }} def _wire(res: dict) -> dict: """Strip non-serializable timing object down to what the UI needs.""" a = res["analysis"] return { "id": res["finding"]["id"], "rule": res["finding"].get("rule", ""), "scanner": res["finding"].get("scanner", ""), "message": res["finding"].get("message", ""), "path": res["finding"].get("path", ""), "category": a.get("category", ""), "cvss": a.get("severity_cvss", 0), "label": a.get("severity_label", "info"), "exploitability": a.get("exploitability", ""), "fp": a.get("false_positive_likelihood", 0), "rationale": a.get("rationale", ""), "remediation": res.get("remediation"), "ttft_ms": round(res["timing"].ttft_ms, 1) if res["timing"].ttft_ms else None, }