subhdotsol commited on
Commit
37ab50f
·
1 Parent(s): a3c4cc9

feat(app): add lifespan hook — creates RedTeamEnvironment on startup

Browse files
Files changed (1) hide show
  1. server/app.py +13 -1
server/app.py CHANGED
@@ -1,10 +1,22 @@
 
1
  from fastapi import FastAPI
2
  from fastapi.middleware.cors import CORSMiddleware
 
 
 
 
 
 
 
 
 
 
 
3
 
4
  app = FastAPI(
5
  title = "RedTeamOS",
6
- description = "AI Red-Teaming Environment for Safety Research",
7
  version = "0.1.0",
 
8
  )
9
 
10
  app.add_middleware(
 
1
+ from contextlib import asynccontextmanager
2
  from fastapi import FastAPI
3
  from fastapi.middleware.cors import CORSMiddleware
4
+ from server.environment import RedTeamEnvironment
5
+ from server.config import get_settings
6
+
7
+ env: RedTeamEnvironment = None
8
+
9
+ @asynccontextmanager
10
+ async def lifespan(app: FastAPI):
11
+ global env
12
+ settings = get_settings()
13
+ env = RedTeamEnvironment(max_turns=settings.max_turns)
14
+ yield
15
 
16
  app = FastAPI(
17
  title = "RedTeamOS",
 
18
  version = "0.1.0",
19
+ lifespan = lifespan,
20
  )
21
 
22
  app.add_middleware(