server/app.py

from contextlib import asynccontextmanager
from pathlib import Path
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from models import AttackAction, StepResult, ResetResponse, EpisodeState, AutoAttackRequest
from server.environment import RedTeamEnvironment
from server.config import get_settings

from rewards.compute_rewards import RewardComputer
from llm.pipeline import run_llm_pipeline, reset_conversation
from llm.automated_attacker import generate_automated_attack

env: RedTeamEnvironment = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global env
    settings = get_settings()
    env = RedTeamEnvironment(max_turns=settings.max_turns)

    reward_computer = RewardComputer()
    env.set_reward_computer(reward_computer)
    env.set_llm_pipeline(run_llm_pipeline)
    yield

app = FastAPI(
    title       = "BreachOS",
    version     = "0.1.0",
    lifespan    = lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins     = ["*"],
    allow_credentials = True,
    allow_methods     = ["*"],
    allow_headers     = ["*"],
)

_FRONTEND = Path(__file__).parent.parent / "frontend"
if _FRONTEND.exists():
    app.mount("/static", StaticFiles(directory=str(_FRONTEND)), name="static")

@app.get("/", include_in_schema=False)
async def serve_ui():
    return FileResponse(str(_FRONTEND / "index.html"))

@app.get("/health")
async def health_check():
    return {"status": "healthy", "version": "0.1.0"}

@app.post("/reset", response_model=ResetResponse)
async def reset_episode():
    try:
        reset_conversation()
        observation = await env.reset()
        return ResetResponse(observation=observation, episode_id=observation.episode_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/step", response_model=StepResult)
async def step_episode(action: AttackAction):
    try:
        result = await env.step(action)
        return result
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/state", response_model=EpisodeState)
async def get_state():
    return env.get_state()

@app.get("/history")
async def get_history():
    return {"history": env.get_history()}

@app.post("/grade")
async def grade_episode():
    if env.is_active:
        raise HTTPException(status_code=400, detail="Episode still active — finish it before grading.")

    history = env.get_history()
    if not history:
        raise HTTPException(status_code=400, detail="No episode history to grade.")

    from graders.programmatic_grader import grade_episode as do_grade
    result = do_grade(history)
    result["episode_id"] = env.episode_id
    return result

@app.post("/auto-attack")
async def auto_attack(request: AutoAttackRequest):
    if not env.is_active:
        raise HTTPException(status_code=400, detail="No active episode.")
    framing = generate_automated_attack(request.strategy_type.value, request.target_category.value)
    return {"framing": framing}


def main():
    import uvicorn
    uvicorn.run("server.app:app", host="0.0.0.0", port=7860, reload=False)


if __name__ == "__main__":
    main()