subhdotsol commited on
Commit
db1fcf1
·
1 Parent(s): 4ec98de

feat(environment): add get_state() — synchronous EpisodeState snapshot

Browse files
Files changed (1) hide show
  1. server/environment.py +10 -1
server/environment.py CHANGED
@@ -69,7 +69,7 @@ class RedTeamEnvironment:
69
  for h in self.attack_history[-5:]
70
  ]
71
 
72
- # ------ Step 3: Call Person 3'LLM pipeline ------
73
  if self.llm_pipeline:
74
  llm_result = await self.llm_pipeline(action, conversation_history)
75
  else:
@@ -145,3 +145,12 @@ class RedTeamEnvironment:
145
  observation = observation,
146
  reward = reward_result.get("total_reward", 0.0),
147
  )
 
 
 
 
 
 
 
 
 
 
69
  for h in self.attack_history[-5:]
70
  ]
71
 
72
+ # ------ Step 3: Call Person 3's LLM pipeline ------
73
  if self.llm_pipeline:
74
  llm_result = await self.llm_pipeline(action, conversation_history)
75
  else:
 
145
  observation = observation,
146
  reward = reward_result.get("total_reward", 0.0),
147
  )
148
+
149
+ def get_state(self) -> EpisodeState:
150
+ return EpisodeState(
151
+ episode_id = self.episode_id or "none",
152
+ turn = self.turn,
153
+ max_turns = self.max_turns,
154
+ attacks_so_far = len(self.attack_history),
155
+ is_active = self.is_active,
156
+ )