backend/main.py

"""
FlashTriage — entry point.

Run:
    cp .env.example .env   # add your CEREBRAS_API_KEY
    pip install -r requirements.txt
    uvicorn backend.main:app --reload --port 8000
Then open http://localhost:8000
"""
from __future__ import annotations

import json
import os
from pathlib import Path

import httpx
from dotenv import load_dotenv
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse

from .agents import Swarm
from .cerebras_client import CerebrasClient, MockClient
from .models import normalize

load_dotenv(Path(__file__).resolve().parent.parent / ".env")

ROOT = Path(__file__).resolve().parent.parent
app = FastAPI(title="FlashTriage")


def _cerebras():
    key = os.getenv("CEREBRAS_API_KEY", "").strip()
    if os.getenv("USE_MOCK", "0") == "1" or not key:
        return MockClient(provider="cerebras"), "cerebras (mock)"
    return CerebrasClient(
        api_key=key,
        base_url=os.getenv("CEREBRAS_BASE_URL", "https://api.cerebras.ai/v1"),
        model=os.getenv("CEREBRAS_MODEL", "gemma-4-31b"),
        provider="cerebras",
    ), os.getenv("CEREBRAS_MODEL", "gemma-4-31b")


def _baseline():
    """Real second provider if configured; else None (caller falls back to sequential).
    Point this at Ollama on your RTX 5090 for a real local-GPU-vs-Cerebras race:
    BASELINE_BASE_URL=http://localhost:11434/api  BASELINE_API_KEY=ollama  BASELINE_MODEL=gemma5090:latest
    """
    url = os.getenv("BASELINE_BASE_URL", "").strip()
    key = os.getenv("BASELINE_API_KEY", "").strip()
    model = os.getenv("BASELINE_MODEL", "").strip()
    if url and model:
        label = os.getenv("BASELINE_LABEL", "").strip() or model
        return CerebrasClient(api_key=key or "none", base_url=url, model=model,
                      provider="baseline", profile="ollama_native"), label
    return None, None


def _swarm(client):
    return Swarm(
        client,
        max_concurrency=int(os.getenv("MAX_CONCURRENCY", "10")),
        remediate_min_cvss=float(os.getenv("REMEDIATE_MIN_CVSS", "7.0")),
    )


def _sse(obj: dict) -> str:
    return f"data: {json.dumps(obj)}\n\n"


@app.get("/", response_class=HTMLResponse)
def index():
    return (ROOT / "frontend" / "index.html").read_text(encoding="utf-8")


@app.get("/api/sample")
def sample():
    return JSONResponse(json.loads((ROOT / "samples" / "sample_findings.json").read_text()))


@app.post("/api/triage")
async def triage(req: Request):
    """
    Body: { "payload": <scanner output: SARIF|JSON|text|list>,
            "lane": "cerebras" | "baseline",
            "limit": int }
    Streams SSE: meta -> finding* -> rollup -> summary
    """
    body = await req.json()
    findings = normalize(body.get("payload"))
    limit = int(body.get("limit") or 0)
    if limit > 0:
        findings = findings[:limit]
    lane = body.get("lane", "cerebras")

    if lane == "baseline":
        bclient, blabel = _baseline()
        if bclient is not None:                       # real second provider, sequential
            client, model_label, parallel = bclient, blabel, False
        else:                                         # honest fallback: same model, sequential
            client, model_label, parallel = _cerebras()[0], _cerebras()[1] + " (sequential)", False
    else:
        client, model_label = _cerebras()
        parallel = True

    swarm = _swarm(client)

    async def gen():
        yield _sse({"type": "meta", "data": {
            "lane": lane, "model": model_label,
            "parallel": parallel, "count": len(findings)}})
        async for ev in swarm.run_batch(findings, parallel=parallel):
            yield _sse(ev)

    return StreamingResponse(gen(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


@app.post("/api/deepdive")
async def deepdive(req: Request):
    """
    Body: { "context": str, "image_base64": str (no data: prefix), "media_type": str }
    Multimodal root-cause analysis on a single finding's screenshot.
    """
    body = await req.json()
    client, model_label = _cerebras()
    swarm = _swarm(client)
    async with httpx.AsyncClient() as http:
        out = await swarm.vision_rca(
            http,
            context=body.get("context", "Analyze this security finding."),
            image_b64=body.get("image_base64", ""),
            media_type=body.get("media_type", "image/png"),
        )
    return JSONResponse({"model": model_label, "rca": out["rca"],
                         "ttft_ms": out["timing"].ttft_ms,
                         "wall_ms": out["timing"].wall_ms})