File size: 3,613 Bytes
da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba da0a862 ee929ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | from __future__ import annotations
import argparse
from pathlib import Path
from typing import Annotated, Any
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, Field
from server.environment import SecretsAuditEnvironment
app = FastAPI(title="SecretsAuditEnv", version="4.0")
environment = SecretsAuditEnvironment(repo_root=Path(__file__).resolve().parents[1])
class ResetRequest(BaseModel):
task_id: int | str | None = None
id: int | str | None = None
difficulty: str | None = None
class StepRequest(BaseModel):
action: str = Field(..., min_length=1)
def _parse_task_id(raw: Any) -> int:
if raw is None:
return 1
if isinstance(raw, int):
task_id = raw
else:
text = str(raw).strip()
if text.startswith("task_"):
text = text.split("_", 1)[1]
try:
task_id = int(text)
except ValueError as exc:
raise HTTPException(status_code=422, detail="Invalid task_id value.") from exc
if not 1 <= task_id <= 13:
raise HTTPException(status_code=422, detail="task_id must be between 1 and 13.")
return task_id
_DIFFICULTY_TO_TASK_ID = {"easy": 1, "medium": 6, "hard": 11}
@app.get("/")
def root() -> dict:
return {
"name": "SecretsAuditEnv",
"version": "4.0",
"status": "ok",
"docs": "/docs",
"web_ui": "/web",
"tasks_endpoint": "/tasks",
}
@app.get("/health")
def health() -> dict:
return {"status": "ok"}
@app.get("/state")
def get_state() -> dict:
return environment.state()
@app.get("/tasks")
def get_tasks() -> dict:
return {"tasks": environment.list_tasks()}
@app.post("/reset")
def reset(
request: ResetRequest | None = None,
task_id: Annotated[str | int | None, Query()] = None,
) -> dict:
try:
raw_task_id: Any = task_id
if raw_task_id is None and request is not None:
if request.difficulty and request.task_id is None and request.id is None:
raw_task_id = _DIFFICULTY_TO_TASK_ID.get(request.difficulty, 1)
else:
raw_task_id = request.task_id if request.task_id is not None else request.id
resolved_task_id = _parse_task_id(raw_task_id)
return environment.reset(resolved_task_id)
except FileNotFoundError as exc:
raise HTTPException(status_code=404, detail=str(exc)) from exc
@app.post("/step")
def step(request: StepRequest) -> dict:
try:
state = environment.step(request.action)
session = (state.get("session") or {})
return {
**state,
"reward": session.get("reward", 0.0),
"cumulative_reward": session.get("reward", 0.0),
"done": float(session.get("reward", 0.0)) >= 1.0,
"step": session.get("step_count", 0),
"observation": session.get("observation", ""),
}
except RuntimeError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
@app.get("/web", response_class=HTMLResponse)
async def web_ui() -> str:
html_path = Path(__file__).parent / "web_ui.html"
return html_path.read_text(encoding="utf-8")
def main() -> None:
parser = argparse.ArgumentParser(description="Run the SecretsAuditEnv server.")
parser.add_argument("--host", default="0.0.0.0")
parser.add_argument("--port", type=int, default=7860)
args = parser.parse_args()
import uvicorn
uvicorn.run(app, host=args.host, port=args.port)
if __name__ == "__main__":
main()
|