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

feat(app): wire RewardComputer and run_llm_pipeline into environment on startup

Browse files
Files changed (1) hide show
  1. server/app.py +7 -0
server/app.py CHANGED
@@ -4,6 +4,9 @@ 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
@@ -11,6 +14,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(
 
4
  from server.environment import RedTeamEnvironment
5
  from server.config import get_settings
6
 
7
+ from rewards.compute_rewards import RewardComputer
8
+ from llm.pipeline import run_llm_pipeline
9
+
10
  env: RedTeamEnvironment = None
11
 
12
  @asynccontextmanager
 
14
  global env
15
  settings = get_settings()
16
  env = RedTeamEnvironment(max_turns=settings.max_turns)
17
+
18
+ reward_computer = RewardComputer()
19
+ env.set_reward_computer(reward_computer)
20
+ env.set_llm_pipeline(run_llm_pipeline)
21
  yield
22
 
23
  app = FastAPI(