Spaces:
Sleeping
Sleeping
Commit ·
b8bc48b
1
Parent(s): d44c135
feat(environment): scaffold RedTeamEnvironment class with state fields
Browse files- server/environment.py +25 -0
server/environment.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import uuid
|
| 2 |
+
import logging
|
| 3 |
+
from typing import Optional, Callable
|
| 4 |
+
from datetime import datetime
|
| 5 |
+
|
| 6 |
+
from models import (
|
| 7 |
+
AttackAction,
|
| 8 |
+
RedTeamObservation,
|
| 9 |
+
EpisodeState,
|
| 10 |
+
StepResult,
|
| 11 |
+
)
|
| 12 |
+
|
| 13 |
+
logger = logging.getLogger(__name__)
|
| 14 |
+
|
| 15 |
+
class RedTeamEnvironment:
|
| 16 |
+
def __init__(self, max_turns: int = 10):
|
| 17 |
+
self.max_turns: int = max_turns
|
| 18 |
+
self.turn: int = 0
|
| 19 |
+
self.attack_history: list[dict] = []
|
| 20 |
+
self.episode_id: Optional[str] = None
|
| 21 |
+
self.is_active: bool = False
|
| 22 |
+
self.created_at: Optional[datetime] = None
|
| 23 |
+
|
| 24 |
+
self.reward_computer: Optional[Callable] = None
|
| 25 |
+
self.llm_pipeline: Optional[Callable] = None
|