Spaces:
Sleeping
Sleeping
| """ | |
| Backlog-scale triage via the Cerebras Batch API (Private Preview). | |
| The live console is for interactive, real-time triage. This is the other half of the | |
| enterprise story: point it at a *huge* backlog (up to 50k findings) and let Cerebras | |
| process it asynchronously, guaranteed within 24h. Same Analyst agent + strict schema. | |
| python -m backend.batch_triage huge_scan.json | |
| Notes: | |
| - Batch is Private Preview; your model must be batch-enabled for your org. | |
| - Minimum 10 requests per batch; this script auto-chunks at 50k. | |
| - Results are NOT ordered — we match on custom_id. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import sys | |
| import time | |
| import httpx | |
| from dotenv import load_dotenv | |
| from .agents import ANALYST_SCHEMA, ANALYST_SYS, _analyst_user | |
| from .models import normalize | |
| load_dotenv() | |
| BASE = os.getenv("CEREBRAS_BASE_URL", "https://api.cerebras.ai/v1").rstrip("/") | |
| MODEL = os.getenv("CEREBRAS_MODEL", "gemma-4-31b") | |
| def _headers() -> dict: | |
| key = os.environ.get("CEREBRAS_API_KEY") | |
| if not key: | |
| sys.exit("Set CEREBRAS_API_KEY to run a batch.") | |
| return {"Authorization": f"Bearer {key}"} | |
| def build_jsonl(findings: list[dict], path: str) -> None: | |
| with open(path, "w", encoding="utf-8", newline="\n") as f: | |
| for i, fd in enumerate(findings): | |
| req = { | |
| "custom_id": fd["id"] or f"f-{i}", | |
| "method": "POST", | |
| "url": "/v1/chat/completions", | |
| "body": { | |
| "model": MODEL, | |
| "max_completion_tokens": 350, | |
| "messages": [ | |
| {"role": "system", "content": ANALYST_SYS}, | |
| {"role": "user", "content": _analyst_user(fd)}, | |
| ], | |
| "response_format": {"type": "json_schema", "json_schema": | |
| {"name": "out", "strict": True, "schema": ANALYST_SCHEMA}}, | |
| }, | |
| } | |
| f.write(json.dumps(req) + "\n") | |
| def main(src: str) -> None: | |
| findings = normalize(open(src, encoding="utf-8").read()) | |
| if len(findings) < 10: | |
| sys.exit("Batch needs >= 10 findings; use the live console for small scans.") | |
| print(f"{len(findings)} findings -> batch") | |
| build_jsonl(findings, "/tmp/batch_in.jsonl") | |
| with httpx.Client(timeout=120) as c: | |
| up = c.post(f"{BASE}/files", headers=_headers(), | |
| files={"purpose": (None, "batch"), | |
| "file": ("batch_in.jsonl", open("/tmp/batch_in.jsonl", "rb"))}) | |
| up.raise_for_status() | |
| file_id = up.json()["id"] | |
| print("uploaded:", file_id) | |
| b = c.post(f"{BASE}/batches", headers={**_headers(), "Content-Type": "application/json"}, | |
| json={"input_file_id": file_id, "endpoint": "/v1/chat/completions", | |
| "completion_window": "24h", | |
| "metadata": {"app": "flashtriage"}}) | |
| b.raise_for_status() | |
| bid = b.json()["id"] | |
| print("batch:", bid) | |
| while True: | |
| st = c.get(f"{BASE}/batches/{bid}", headers=_headers()).json() | |
| rc = st.get("request_counts", {}) | |
| print(f" {st['status']} {rc.get('completed',0)}/{rc.get('total',0)}") | |
| if st["status"] in ("completed", "failed", "expired", "cancelled"): | |
| break | |
| time.sleep(10) | |
| out_id = st.get("output_file_id") | |
| if not out_id: | |
| sys.exit(f"no output ({st['status']})") | |
| res = c.get(f"{BASE}/files/{out_id}/content", headers=_headers()) | |
| open("batch_results.jsonl", "wb").write(res.content) | |
| print("wrote batch_results.jsonl") | |
| if __name__ == "__main__": | |
| main(sys.argv[1] if len(sys.argv) > 1 else "samples/sample_findings.json") | |